Mercurial > repos > pointsource_pde / changeset
changeset
Initial working version for known convectivity and diffusivity
Thu, 26 Feb 2026 09:32:12 -0500
- author
- Tuomo Valkonen <tuomov@iki.fi>
- date
- Thu, 26 Feb 2026 09:32:12 -0500
- changeset 1
- a4137aedcb3a
- parent 0
- 7ec1cfe19a24
- child 2
- 69002abe5dcb
- child 3
- c3a4f4bb87f7
Initial working version for known convectivity and diffusivity
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/.clang-format Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,15 @@ +--- +BasedOnStyle: LLVM +IndentWidth: 4 +ColumnLimit: 98 +--- +Language: Cpp +# Force pointers to the type for C++. +DerivePointerAlignment: false +PointerAlignment: Left +NamespaceIndentation: All +AllowShortIfStatementsOnASingleLine: false +AllowShortLoopsOnASingleLine: false +AllowShortBlocksOnASingleLine: Never +AllowShortFunctionsOnASingleLine: None +---
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/.gitignore Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,1 @@ +.hgignore \ No newline at end of file
--- a/.hgignore Tue Apr 08 13:08:16 2025 -0500 +++ b/.hgignore Thu Feb 26 09:32:12 2026 -0500 @@ -1,6 +1,12 @@ -^target/ -^debug_out/ -^pointsource.._.*\.txt -flamegraph.svg -DEADJOE -.*\.orig +syntax: glob +target/ +Cargo.lock +.clangd +.envrc +.zed/settings.json +res/ +*.npz +__pycache__ +*.json +*.orig +*.png
--- a/Cargo.toml Tue Apr 08 13:08:16 2025 -0500 +++ b/Cargo.toml Thu Feb 26 09:32:12 2026 -0500 @@ -12,18 +12,47 @@ categories = ["mathematics", "science", "computer-vision"] [dependencies.alg_tools] -version = "~0.3.1" +version = "~0.4.0-dev" path = "../alg_tools" default-features = false -features = ["nightly"] +features = ["nightly", "pyo3"] + +[dependencies.pointsource_algs] +version = "~3.0.0-dev" +path = "../pointsource_algs" + +[dependencies.measures] +version = "~0.1.0" +path = "../measures" +features = ["pyo3"] [dependencies] -serde = { version = "1.0", features = ["derive"] } -cxx = "1.0" +serde = { version = "~1.0", features = ["derive"] } +cxx = "~1.0" +pyo3 = { version = "0.27.1", features = ["anyhow"] } +anyhow = "~1.0.95" +itertools = "~0.14.0" +dolfinx-sys = { path = "dolfinx-sys" } +nanobind-sys = { path = "nanobind-sys" } +clap = "~4.5.38" +serde_with = "~3.12.0" +serde_json = "1.0.140" +nalgebra = "0.34" +numpy = { version = "0.27.0", features = ["nalgebra"] } +ndarray = "0.16.1" +log = { version = "0.4.28", features = [ + "max_level_debug", + "release_max_level_warn", +] } +colog = "1.3.0" [build-dependencies] -cxx-build = "1.0" +build-print = "~0.1.1" +cc = { version = "~1.2.22", features = ["parallel"] } +cxx-build = "~1.0" +#pyo3 = "~0.25.0" regex = "~1.11.0" +conda-build = { path = "conda-build" } [profile.release] debug = true
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/README.md Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,96 @@ + +# Pointsource-PDE + +## Prerequisites + +## Installation and usage + +### Installing dependencies + +Python dependencies are managed Conda and by the Cargo build system of [Rust]. +On some platforms, you can use alternative methods. + +#### Phase 1: Python + +##### Most platforms: + +- Fenicsx installed in Conda according to instructions +- SciFEM and SciPY (`conda install conda-forge::scifem conda-forge::scipy`) + +##### Debian/Ubuntu + +You may br able to use the system package manager instead of Conda, but beware of obsolete versions. + +#### Phase 2: Rust + +You will only need to install the “nightly” Rust compiler and the +[GNU Scientific Library] manually. At the time of writing this README, +[alg_tools] also needs to be downloaded separately. + +1. Install the [Rust] infrastructure (including Cargo) with [rustup]. +2. Install a “nightly” release of the Rust compiler. With rustup, installed in + the previous step, this can be done with + ```console + rustup toolchain install nightly + ``` +3. Install [GNU Scientific Library]. On a Mac with [Homebrew] installed, + this can be done with + ```console + brew install gsl + ``` + For other operating systems, suggestions are available in the [rust-GSL] + crate documentation. If not correctly installed, you may need to pass + extra `RUSTFLAGS` options to Cargo in the following steps to locate the + library. + +4. Download [alg_tools] and unpack it under the same directory as this + package. + + [rustup]: https://rustup.rs + [alg_tools]: https://tuomov.iki.fi/software/alg_tools/ + [Rust]: https://www.rust-lang.org/ + [GNU Scientific Library]: https://www.gnu.org/software/gsl/ + [rust-GSL]: https://docs.rs/GSL/6.0.0/rgsl/ + [Homebrew]: https://brew.sh + [arXiv:2212.02991]: https://arxiv.org/abs/2212.02991 + [arXiv:2502.12417]: https://arxiv.org/abs/2502.12417 + [doi:10.46298/jnsao-2023-10433]: http://doi.org/10.46298/jnsao-2023-10433 + +### Building and running the experiments + +To compile and install the program, use +```console +cargo install --path=. +``` +When doing this for the first time, several dependencies will be downloaded. +Now you can run the default set of experiments with +``` +pointsource_pde -o results +``` +The `-o results` option tells `pointsource_pde` to write results in the +`results` directory. The option is required. + +Alternatively, you may build and run the program without installing with +```console +cargo run --release -- -o results +``` +The double-dash separates the options for the Cargo build system +and `pointsource_pde`. + +### Documentation + +Use the `--help` option to get an extensive listing of command line options to +customise algorithm parameters and the experiments performed. + +## Internals + +If you are interested in the program internals, the integrated source code +documentation may be built and opened with +```console +cargo doc # build dependency docs +misc/cargo-d --open # build and open KaTeX-aware docs for this crate +``` +The `cargo-d` script ensures that KaTeX mathematics is rendered in the +generated documentation through an ugly workaround. Unfortunately, +`rustdoc`, akin to Rust largely itself, is stuck in 80's 7-bit gringo ASCII +world, and does not support modern markdown features, such as mathematics.
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/build.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,138 @@ +/*! +Build script for `pointsource_pde` +*/ + +//use pyo3::{prepare_freethreaded_python, types::PyModule, PyErr, Python}; +use conda_build::python_config; +use std::env::split_paths; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::process::Command; + +const REQUIRED_PYTHON_MODULES: [&str; 7] = [ + "dolfinx", "ufl", "numpy", "time", "petsc4py", "mpi4py", "json", +]; + +const DOLFINX_ACCESS_CXX_SOURCES: [&str; 3] = + ["nanobind_helpers.cc", "minmax_p2.cc", "function.cc"]; + +const PY_SRC: [&str; 7] = [ + "src/dolfinx_extras.py", + "src/compose.py", + "src/measure.py", + "src/quadratic_dataterm.py", + "src/convection_diffusion.py", + "src/laser_sampling.py", + "src/full_sampling.py", +]; + +fn main() -> Result<(), std::io::Error> { + // Check Python module availability. + // This doesn't necessarily work in a Conda setup. + // prepare_freethreaded_python(); + // Python::with_gil(|py| { + // for m in REQUIRED_PYTHON_MODULES { + // PyModule::import(py, m)?; + // } + // Ok::<_, PyErr>(()) + // })?; + + let (pyc_e, py_is_conda) = python_config(); + let py_path = pyc_e.and_then(|pyc| pyc.exec_prefix_path())?; + let py_exec = py_path.join("bin").join("python3"); + let import_check_code = REQUIRED_PYTHON_MODULES + .map(|m| format!("import {}\n", m)) + .concat(); + let res = Command::new(&py_exec) + .args(["-c", import_check_code.as_ref()]) + .status()?; + if !res.success() { + return Err(std::io::Error::other(format!( + "Failed to load required Python modules (python execuitable {}{})", + py_exec.display(), + if py_is_conda { " from Conda" } else { "" } + ))); + } + + // Set up local-to-crate paths + let cc_mod_name = "dolfinx_access"; + let crate_name = "pointsource_pde"; + let src = PathBuf::from("src"); + let cc_mod_root = src.join(cc_mod_name); + let bridge = src.join(format!("{cc_mod_name}.rs")); + let cc_sources = DOLFINX_ACCESS_CXX_SOURCES.map(|f| cc_mod_root.join(f)); + let mut include_dirs = Vec::from([PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("include")]); + + // Add (an estimate of) dolfinx ldflags. This shouldn't really be our problem to deal with, + // but everything that is not piss, is shit. + add_dep_ldflags("DEP_DOLFINX_LDFLAGS"); + + // Add external include paths + add_dep_includes(&mut include_dirs, "DEP_DOLFINX_INCLUDE"); + add_dep_includes(&mut include_dirs, "DEP_NANOBIND_INCLUDE"); + + // This may be useful to enable sometimes, when there's + // trouble location GSL. + // pkg_config::Config::new().probe("gsl").unwrap(); + + // Build the bridge + cxx_build::bridge(&bridge) + .std("c++20") + .files(&cc_sources) + .includes(&include_dirs) + .compile(crate_name); + + for dep in PY_SRC { + println!("cargo::rerun-if-changed={}", dep); + } + + let header_prefix = PathBuf::from("include").join(cc_mod_name); + let headers = DOLFINX_ACCESS_CXX_SOURCES + .into_iter() + .map(|s| header_prefix.join(s.replace(".cc", ".h"))); + + for dep in cc_sources + .into_iter() + .chain(headers) + .chain(std::iter::once(bridge)) + { + println!("cargo:rerun-if-changed={}", dep.display()); + } + + // Generate .clangd + let mut dot_clangd = File::create(".clangd")?; + write!( + dot_clangd, + "\ +# Automatically generated by build.rs +CompileFlags: + Add: +" + )?; + for i in include_dirs { + writeln!(dot_clangd, " - --include-directory={}", i.display())?; + } + + Ok(()) +} + +/// Add list of include directions in the environment variable `dep` to +/// `include_dirs`. +fn add_dep_includes(include_dirs: &mut Vec<PathBuf>, dep: &str) { + if let Some(include) = std::env::var_os(&dep) { + include_dirs.extend(split_paths(&include).collect::<Vec<_>>()); + } else { + panic!("{dep} unset."); + } +} + +fn add_dep_ldflags(dep: &str) { + if let Some(ldflags) = std::env::var_os(&dep) { + for lp in split_paths(&ldflags) { + println!("cargo:rustc-link-arg=-Wl,{}", lp.display()); + } + } else { + panic!("{dep} unset."); + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/conda-build/Cargo.toml Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,14 @@ +[package] +name = "conda-build" +version = "0.1.0" +edition = "2021" +authors = ["Tuomo Valkonen <tuomov@iki.fi>"] +description = "Helpers for build.rs scripts for working with external deps managed by conda" + +[lib] +name = "conda_build" +crate-type = ["dylib", "rlib"] + +[dependencies] +anyhow = "~1.0.98" +python-config-rs = "0.1.2"
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/conda-build/src/lib.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,77 @@ +//use anyhow::bail; +//use regex::Regex; +pub use python_config::{Error as PyCError, PythonConfig}; +use std::env; +use std::ffi::OsString; +use std::path::PathBuf; +//use std::process::Command; + +/// Adds CONDA_PREFIX/lib/pkgconfig/ to PKG_CONFIG_PATH +pub fn pkgconfig_conda() -> Result<Option<PathBuf>, anyhow::Error> { + if let Some(conda_root_) = option_env!("CONDA_PREFIX") { + let conda_root = PathBuf::from(conda_root_); + let conda_pkgconfig = conda_root.join("lib/pkgconfig/"); + + // Prepend CONDA_PREFIX to any existing PKG_CONFIG_PATH + let new_path = match option_env!("PKG_CONFIG_PATH") { + Some(path) => { + let mut paths = env::split_paths(&path).collect::<Vec<_>>(); + paths.insert(0, conda_pkgconfig); + env::join_paths(paths)? + } + None => conda_pkgconfig.into(), + }; + + // Set the environment variable + unsafe { + env::set_var("PKG_CONFIG_PATH", &new_path); + } + + Ok(Some(conda_root)) + } else { + Ok(None) + } +} + +/// Looks up conda Python +pub fn python_conda() -> Option<OsString> { + option_env!("CONDA_PREFIX").and_then(|conda_prefix| { + /* This does not point to the right place. + option_env!("CONDA_PYTHON_EXE") + .map(OsString::from) + .or_else(|| { */ + let py = PathBuf::from(conda_prefix).join("bin/python"); + if py.exists() { + Some(py.into()) + } else { + None + } + /* })*/ + }) +} + +/// Returns a the python configuration and indication whether it's from Conda +pub fn python_config() -> (Result<PythonConfig, PyCError>, bool) { + if let Some(py) = python_conda() { + (PythonConfig::interpreter(py), true) + } else { + (Ok(PythonConfig::new()), false) + } +} + +/* +/// Looks up python includes given executable +pub fn python_includes(mut py: OsString) -> Result<Vec<String>, anyhow::Error> { + py.push("3-config"); + let output = Command::new(&py).args(["--includes"]).output()?; + if !output.status.success() { + bail!("Failed to run {}.", py.display()) + } + let re = Regex::new(r"\w+-I\w*").unwrap(); + let includes = re + .split(format!(" {}", String::from_utf8(output.stdout)?).as_str()) + .map(String::from) + .collect(); + Ok(includes) +} +*/
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/dolfinx-sys/Cargo.toml Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,18 @@ +[package] +name = "dolfinx-sys" +version = "0.1.0" +edition = "2021" +authors = ["Tuomo Valkonen <tuomov@iki.fi>"] +description = "Dummy crate for location and linking against system dolfinx library" +links = "dolfinx" + +[build-dependencies] +anyhow = "~1.0.98" +build-print = "~0.1.1" +itertools = "~0.14.0" +pkg-config = "~0.3" +conda-build = { path = "../conda-build" } + +[lib] +name = "dolfinx_sys" +crate-type = ["dylib", "rlib"]
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/dolfinx-sys/build.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,51 @@ +use build_print::info; +use conda_build::pkgconfig_conda; +use itertools::Itertools; +use std::env; + +fn main() -> Result<(), anyhow::Error> { + if let Some(conda_root) = pkgconfig_conda()? { + info!( + "Directing pkg-config to look for fenics-dolfinx within the Conda environment {}.", + conda_root.display() + ); + } + + println!("cargo:rerun-if-changed=build.rs"); + + let mut pkg = pkg_config::Config::new(); + locate_dolfinx(&mut pkg) +} + +fn locate_dolfinx(pkg: &mut pkg_config::Config) -> Result<(), anyhow::Error> { + let lib = pkg.probe("dolfinx").or_else(|e| { + // Debian/Ubuntu has separate real and complex dolfinx. + info!("Could not find dolfinx, trying dolfinx_real (as the library is named if installed through a Debian/Ubuntu package)."); + // If that's also not found, return original error instead of Debian/Ubuntu-specic error. + pkg.probe("dolfinx_real").map_err(|_| e) + })?; + + // We also need fmt for linking to work. + pkg.probe("fmt")?; + + let cleaned_paths = lib.include_paths.iter().unique().collect::<Vec<_>>(); + let includes = env::join_paths(cleaned_paths)?.into_string().unwrap(); + println!("cargo:include={}", includes); + + println!( + "cargo:ldflags={}", + env::join_paths(lib.ld_args.iter().map(|s| s.join(",")))? + .into_string() + .unwrap() + ); + + // let cleaned_rpaths = lib.link_paths.iter().unique().collect::<Vec<_>>(); + // let rpaths = env::join_paths(cleaned_rpaths)?.into_string().unwrap(); + // println!("cargo:rpath={}", rpaths); + + // for lp in lib.link_paths.iter().unique() { + // println!("cargo:rustc-link-arg=-Wl,-rpath,{}", lp.display()); + // } + + Ok(()) +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/dolfinx-sys/src/lib.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,4 @@ +/*! +This is a dummy crate to help linking against system dolfinx / fenics libraries. +It contains no functions. +*/
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/experiments/laser_and_mirrors.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,38 @@ +# +# Laser and mirrors convection-diffusion experiment with known diffusivity and convectivity +# +import laser_and_mirrors_aux +from measures import DiscreteMeasure_2_f64 +from pointsource_pde import Problem, RegTerm +from pointsource_pde.compose import ( + ComposeFnWithOperator, + InjectSecond, +) + +# Give name to the problem +name = "laser_and_mirrors" + +# Override algorithm settings +algorithm_overrides = laser_and_mirrors_aux.algorithm_overrides + + +# Setup routine +def setup(prefix): + dat, auxtrue, μ_bound, μ_true, plot_factory = laser_and_mirrors_aux.generic_setup( + prefix + ) + + μ0 = DiscreteMeasure_2_f64([]) + + reg = RegTerm.NonnegRadon( + 10e-7 # 5e-7 # 3e-7 + ) + + dat_simple = ComposeFnWithOperator(dat, InjectSecond(auxtrue), xbound=μ_bound) + dat_simple.curvature_bound_components = lambda: (None, None) + print("Initial value:", dat_simple.apply(μ0)) + print("Value at true μ:", dat_simple.apply(μ_true)) + + return Problem( + dataterm=dat_simple, regterm=reg, μinit=μ0, plot_factory=plot_factory + )
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/experiments/laser_and_mirrors_aux.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,196 @@ +# +# Laser and mirrors convection-diffusion experiment with unknown diffusivity and convectivity; +# Includes `generic_setup` for the known case as well. +# + +import numpy as np +from dolfinx import fem +from measures import DiscreteMeasure_2_f64 +from numpy.linalg import norm +from pointsource_pde import Problem, RegTerm +from pointsource_pde.compose import ( + ComposeFnWithOperator, + SumOfSeparableFunctions, +) +from pointsource_pde.convection_diffusion import ( + ConvectionDiffusion, + PlotFactory, + QuadraticRegularisation, + XBound, +) +from pointsource_pde.laser_sampling import LaserSampling +from pointsource_pde.quadratic_dataterm import QuadraticDataTerm + +# Give name to the problem +name = "laser_and_mirrors_aux" + +# Fine-tune algorithms +alg_finetune = { + "tau0": 0.99, + "theta0": 0.01, + "inner_method": "PP", + "inner_pdps_τσ0": (19.8, 0.05), + "inner_pp_τ": [1000.0, 1000.0], + "merge": True, + "fitness_merging": True, +} + +algorithm_overrides = { + "RadonFB": alg_finetune, + "RadonSlidingFB": alg_finetune, + "RadonPDPS": alg_finetune, + "RadonSlidingPDPS": alg_finetune, +} + + +# Setup shared between experiment with known/unknown conductivity and diffusivity +def generic_setup(prefix, simple=True): + + # Create true source + μ_true = DiscreteMeasure_2_f64( + [([0.1, 0.3], 0.08), ([0.4, 0.25], 0.05), ([0.25, 0.13], 0.06)] + ) + + # Create convection-diffusion PDE + pde = ConvectionDiffusion( + u0=lambda x: 0.0 * x[1], + g=lambda x: 0.0 * x[1], + domain_size=[0, 0.5, 0, 0.5], + num_steps=50, + ) + + # Set up lasers and mirrors + lasers = [ + np.array([0.1, 0.1]), + np.array([0.1, 0.4]), + np.array([0.4, 0.1]), + np.array([0.4, 0.1]), + ] + mirrors_per_side = 10 + x0, x1, y0, y1 = ( + pde.domain_size[0], + pde.domain_size[1], + pde.domain_size[2], + pde.domain_size[3], + ) + beams = [ + (src, np.array(dst)) + for src in lasers + for i in range(0, mirrors_per_side) + for z in [(i + 1) / (mirrors_per_side + 1)] + for dst in [ + [z * (x1 - x0) + x0, y0], + [z * (x1 - x0) + x0, y1], + [x0, z * (y1 - y0) + y0], + [x1, z * (y1 - y0) + y0], + ] + ] + L_i = LaserSampling(pde.V, pde.domain, pde.nx, pde.ny, beams) + + v0 = 1.0 + + if simple: + # Scalar convectivity and diffusivity + r = 0.5 + θ = 30 * np.pi / 180 + c1 = r * np.cos(θ) + c2 = r * np.sin(θ) + k = 0.01 # diffusive + # k = 0.001 # high wind + else: + # Convectivity and diffusivity fields + def velocity_x(x): + r = np.sqrt(x[0] ** 2 + x[1] ** 2 + 1e-10) + return v0 * x[0] / r + + def velocity_y(x): + r = np.sqrt(x[0] ** 2 + x[1] ** 2 + 1e-10) + return v0 * x[1] / r + + c1 = fem.Function(pde.V) + c1.interpolate(velocity_x) + c2 = fem.Function(pde.V) + c2.interpolate(velocity_y) + k = fem.Function(pde.V) + k0 = 1.0 + k.interpolate(lambda x: k0 + 0.2 * np.sin(5 * x[0]) * np.cos(5 * x[1])) # k is + + auxtrue = (k, (c1, c2)) + + # Simulate measurements + u_n_list = pde.apply((μ_true, (k, (c1, c2)))) + b_true = [L_i.apply(u) for u in u_n_list] + b = [b_i + L_i.noise(0.01 * norm(b_i) / np.sqrt(np.size(b_i))) for b_i in b_true] + + # Estimate bounds on domains and values + dat_bound = 10 * max([norm(b_i) for b_i in b]) + μ_bound = 1.3 * sum([alpha for _point, alpha in μ_true.iter_padded()]) + xbound = XBound(mu_dual=μ_bound, k_min=k, k_max=k, c_Linf=r) + + # Create data term + dat = ComposeFnWithOperator( + SumOfSeparableFunctions( + [ + QuadraticDataTerm( + L_i, + data, + xbound=dat_bound, + b_norm=norm(b), + λ=pde.dt / len(beams), + ) + for data in b + ] + ), + pde, + xbound=[xbound], + xbound_pair=[xbound], + ) + + # Plot true and initial data + plot_factory = PlotFactory(pde) + plotter = plot_factory.produce(prefix) + ω, _ = dat.diff((DiscreteMeasure_2_f64([]), auxtrue)) + ω_true, _ = dat.diff((μ_true, auxtrue)) + np.savez_compressed( + "%s/omega_0.npz" % prefix, + true_mu=[[point[0], point[1], alpha] for point, alpha in μ_true.iter_padded()], + omega0=ω.x.array, + true_omega=ω_true.x.array, + true_u_n_list_array=[u_n_list[i].x.array for i in plotter.plot_frames], + ) + + return dat, auxtrue, μ_bound, μ_true, plot_factory + + +def setup(prefix): + dat, auxtrue, _μ_bound, μ_true, plot_factory = generic_setup(prefix) + + reg = RegTerm.NonnegRadon(10e-7) + + μ0 = DiscreteMeasure_2_f64([]) + aux0 = auxtrue + aux = QuadraticRegularisation(0.001, aux0) + ax = aux.apply(aux0) + inival = dat.apply((μ0, aux0)) + print("Initial value:", inival + ax, ax) + print("Value at true μ:", dat.apply((μ_true, auxtrue)) + ax) + + # auxinit = np.c_[ + # np.zeros_like(k), + # np.zeros_like(c1), + # np.zeros_like(c2), + # ] + # print(np.shape(auxinit)) + + # We need Θ² such that ⟨F'(μ+Δ, aux-F'(μ, aux)|Δ⟩ ≤ Θ²|γ|(c_2), where Δ=(π_♯^1-π_♯^0)γ. + # We have + # + # ⟨F'(μ+Δ, aux)-F'(μ, aux)|Δ⟩ ≤ ‖⟨F'(μ+Δ, aux)-F'(μ, aux)‖ ‖Δ‖ + # ≤ L_{F'(·, aux)} ‖Δ‖², + # + # so Θ² ≥ L_{F'(·, aux)} ‖Δ‖² / |γ|(c_2). Thsi doesn't give anything useful if the + # transport is very small in distance. + + dat.curvature_bound_components = lambda: (None, None) # 1.0 + + return Problem(dat, reg, aux, aux0, μ0, plot_factory=plot_factory)
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/experiments/tests/test_minimisation.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,34 @@ +from dolfinx_access import eval_Function_f64, min_Function_f64_p2 + + +# A simple test that the finite-dimensional subproblem is solved correctly +def test_minimisation(pde): + test = fem.Function(pde.V) + + # Beale function + def beale(x): + return ( + (1.5 - x[0] + x[0] * x[1]) ** 2 + + (2.25 - x[0] + x[0] * (x[1] ** 2)) ** 2 + + (2.625 - x[0] + x[0] * (x[1] ** 3)) ** 2 + ) + + test.interpolate(beale) + m = min_Function_f64_p2(test._cpp_object) + print("Beale minimizer:", m.x, m.v) + + def quadratic_testf(x): + return x[0] ** 2 - 5 * x[1] ** 2 + x[0] * x[1] - x[0] + x[1] + + testf = quadratic_testf + test.interpolate(testf) + n = 256 + dist = 0.0 + for x in np.linspace(pde.domain_size[0], pde.domain_size[1], n): + for y in np.linspace(pde.domain_size[2], pde.domain_size[3], n): + p = [x, y, 0.0] + v = eval_Function_f64(test._cpp_object, p) + v0 = testf(p) + d = np.linalg.norm(v - v0) + dist = max(d, dist) + print("Quadratic function evaluation max discrepancy:", dist)
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/include/dolfinx_access/function.h Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,49 @@ +#pragma once +#include "dolfinx/geometry/BoundingBoxTree.h" +#include <cmath> + +#pragma once +#include "nanobind_helpers.h" + +#pragma once +#include "rust/cxx.h" + +namespace dolfinx_access { + struct FunctionInfo; + + typedef geometry::BoundingBoxTree<double> BBTree_f64; + extern void drop_Function_f64(Function_f64* f); + + extern int32_t cell_Mesh_f64(std::shared_ptr<const mesh::Mesh<double>> mesh, + const std::array<double, 3>& x); + + extern int32_t cell_FunctionSpace_f64(FunctionSpace<double> const* v, + const std::array<double, 3>& x); + + extern int32_t cell_Function_f64(Function_f64 const* f, const std::array<double, 3>& x); + + extern double eval_Function_f64_cell(Function_f64 const* f, const std::array<double, 3>& x, + int32_t cell); + + extern double eval_Function_f64(Function_f64 const* f, const std::array<double, 3>& x); + + extern std::array<double, 6> eval_Function_f64_cell_6(Function_f64 const* f, + const std::array<double, 3 * 6>& x, + int32_t cell_index); + + FunctionInfo info_Function_f64(Function_f64 const* f); + + std::array<double, 3 * 3> cell_coords_Function_f64_triangle(Function_f64 const* f, + int32_t cell); + + rust::Slice<const double> data_Function_f64(Function_f64 const* f); + + rust::Slice<double> data_mut_Function_f64(Function_f64* f); + + void check_compat_Function_f64(Function_f64 const* f, Function_f64 const* g); + + Function_f64* similar_Function_f64(Function_f64 const* f); + + Function_f64* clone_Function_f64(Function_f64 const* f); + +} // namespace dolfinx_access
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/include/dolfinx_access/minmax_p2.h Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,13 @@ +#pragma once +#include "nanobind_helpers.h" + +namespace dolfinx_access { + struct CoordValuePair; + + extern CoordValuePair min_Function_f64_p2(Function_f64 const* f); + extern CoordValuePair max_Function_f64_p2(Function_f64 const* f); + extern CoordValuePair minmax_Function_f64_p2(Function_f64 const* f, bool max); + +} // namespace dolfinx_access + +// extern "C" PyObject* PyInit_dolfinx_access();
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/include/dolfinx_access/nanobind_helpers.h Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,25 @@ +#pragma once +#include <cmath> + +#pragma once +#include <dolfinx/fem/Function.h> + +#pragma once +#include <Python.h> + +namespace dolfinx_access { + using namespace dolfinx::fem; + + using PyObject = ::PyObject; + typedef Function<double> Function_f64; + + const Function_f64* cast_Function_f64(const PyObject* obj); + Function_f64* cast_mut_Function_f64(PyObject* obj); + + bool check_Function_f64(const PyObject* obj); + + PyObject* wrap_Function_f64(Function_f64* f); + +} // namespace dolfinx_access + +extern "C" PyObject* PyInit_dolfinx_access();
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/nanobind-sys/Cargo.toml Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,18 @@ +[package] +name = "nanobind-sys" +version = "0.1.0" +edition = "2021" +authors = ["Tuomo Valkonen <tuomov@iki.fi>"] +description = "Dummy crate for locating, building, and linking against nanobind" +links = "nanobind" + +[build-dependencies] +anyhow = "~1.0.98" +build-print = "~0.1.1" +conda-build = { path = "../conda-build" } +cc = "~1.2.22" +#python-config-rs = "~0.1.2" + +[lib] +name = "nanobind_sys" +crate-type = ["dylib", "rlib"]
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/nanobind-sys/build.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,83 @@ +use anyhow::bail; +use build_print::info; +use cc; +use conda_build::python_config; +use std::env; +use std::path::PathBuf; + +const NANOBIND_SOURCES: [&str; 7] = [ + "common.cpp", + "trampoline.cpp", + "nb_func.cpp", + "nb_type.cpp", + "nb_internals.cpp", + "error.cpp", + "implicit.cpp", +]; + +fn main() -> Result<(), anyhow::Error> { + // Need to build it + let (pyc_e, py_conda) = python_config(); + + // This is very clumsy due to PythonConfig::Error not supporting + // conversion into std::error::Error, just std::io::Error. + let (pyc, mut includes, _py_ldflags) = pyc_e + .and_then(|pyc| { + pyc.include_paths() + .and_then(|ip| pyc.ldflags().map(|ld| (pyc, ip, ld))) + }) + .map_err(|e| anyhow::Error::from(std::io::Error::from(e)))?; + + let nanobind_root = pyc + .prefix_path().ok() + .zip(pyc.semantic_version().ok()) + .and_then(|(prefix, version)| { + info!("{}Python {} found at prefix {}", if py_conda { "Conda " } else { ""}, version, prefix.display()); + let nanobind_root = prefix + .join("lib") + .join(format!("python{}.{}", version.major, version.minor)) + .join("site-packages") + .join("nanobind"); + if !nanobind_root.exists() { + info!("… but no nanobind installed there"); + None + } else { + Some(nanobind_root) + } + }) + .map_or_else(|| { + let last_ditch = "/usr/share/nanobind"; + let nanobind_root = PathBuf::from(last_ditch); + if !nanobind_root.exists() { + bail!("Could not find nanobind in known locations. The recommended installation method is with pip at system Python location, or with Conda into the active environment. Also /usr/share/nanobind was attempted."); + } else { + info!("Found in {last_ditch}."); + Ok(nanobind_root) + } + }, Ok)?; + + let nanobind_src = nanobind_root.join("src"); + + println!("cargo:rerun-if-changed=build.rs"); + + //println!("cargo:rustc-link-arg={}", py_ldflags); + + includes.extend([ + nanobind_root.join("include"), + nanobind_root.join("ext/robin_map/include"), + ]); + + println!( + "cargo:include={}", + env::join_paths(&includes)?.into_string().unwrap() + ); + + cc::Build::new() + .std("c++20") + .files(NANOBIND_SOURCES.map(|f| nanobind_src.join(f))) + .includes(includes) + //.flags(pyc.cflags()?) + .compile("nanobind-sys"); + + Ok(()) +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/nanobind-sys/src/lib.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,4 @@ +/*! +This is a dummy crate to help building and linking against nanobind. +It contains no functions. +*/
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/plot.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,268 @@ +import os +import re +import sys +from itertools import chain +from pathlib import Path +from statistics import median + +import matplotlib as mpl +import matplotlib.cm as cm +import matplotlib.colors as colors +import matplotlib.pyplot as plt +import matplotlib.tri as tri +import numpy as np +from dolfinx import fem + +from src.convection_diffusion import ConvectionDiffusion + + +def find_files(directory): + """ + Return a list of Path objects for files named 'fubar%d.npz' + in the given directory, sorted by the integer %d. + """ + directory = Path(directory) + pattern = re.compile(r"^res_(\d+)\.npz$") + + matches = [] + + for path in directory.iterdir(): + if path.is_file(): + m = pattern.match(path.name) + if m: + number = int(m.group(1)) + matches.append((number, path)) + + # Sort numerically by extracted number + matches.sort(key=lambda x: x[0]) + + return matches + + +def plot(prefix): + quantisation = 32 + + iter_files = find_files(prefix) + + pde = ConvectionDiffusion( + u0=lambda x: 0.0 * x[1], + g=lambda x: 0.0 * x[1], # ((x[1] + 2) / 4) ** 2, # nonzero bcs + # w0=lambda x: np.cos(5 * x[1]), + domain_size=[0, 0.5, 0, 0.5], + # nx=128, + # ny=128, + ) + + # Get coordinates + coords = pde.V.tabulate_dof_coordinates() + x_coords, y_coords = coords[:, 0], coords[:, 1] + triang = tri.Triangulation(x_coords, y_coords) + + m = 5 + fig, (axus, axws, axjs, ax_wpbar_plus) = plt.subplots(4, m, figsize=(15, 4)) + plt.tight_layout() + fig.set_size_inches(13, 8) + fig.subplots_adjust( + left=0.02, right=0.95, bottom=0.02, top=0.93, wspace=0.2, hspace=0.2 + ) + for ax in chain(axus, axws, axjs, ax_wpbar_plus): + ax.set_xticklabels([]) + ax.set_yticklabels([]) + ax.set_aspect("equal") + + mpl.rcParams["axes.labelsize"] = 9 + + norm = colors.Normalize(vmin=0, vmax=1) + sm = cm.ScalarMappable(norm=norm, cmap="viridis") + sm.set_array([]) # REQUIRED (matplotlib quirk) + colorbar_u = fig.colorbar(sm, ax=axus[-1]) + colorbar_w = fig.colorbar(sm, ax=axws[-1]) + colorbar_j = fig.colorbar(sm, ax=axjs[-1]) + + def do_plot(index_and_file, show_true=False): + (iter, file) = index_and_file + + print("🎨 Plotting iteration %d..." % iter) + + data0 = np.load("%s/omega_0.npz" % Path(prefix).parent) + ω0 = data0["omega0"] + true_μ = data0["true_mu"] + # true_ω = data0["true_omega"] + true_u_n_list_array = data0["true_u_n_list_array"] + ω0_min = ω0.real.min() + ω0_max = ω0.real.max() + + for ax in chain(axus, axws, axjs, ax_wpbar_plus): + # ax.clear() + for artist in chain(ax.lines, ax.collections, ax.images): + artist.remove() + + data = np.load(file) + u_n_list_array = data["u_n_list_array"] + w_n_list_array = data["w_n_list_array"] + j_n_list_array = data["j_n_list_array"] + # frames = data["frames"] + times = data["times"] + mu = data["mu"] + wp_bar_array = data["wp_bar_array"] + + u_min = min(u.real.min() for u in u_n_list_array) + u_max = max(u.real.max() for u in u_n_list_array) + if show_true: + u_min = min(u_min, min(u.real.min() for u in true_u_n_list_array)) + u_max = max(u_max, max(u.real.max() for u in true_u_n_list_array)) + + w_min = min( + min(w.real.min() for w in w_n_list_array), wp_bar_array.real.min(), ω0_min + ) + w_max = max( + max(w.real.max() for w in w_n_list_array), wp_bar_array.real.max(), ω0_max + ) + j_min = min(j.real.min() for j in j_n_list_array) + j_max = max(j.real.max() for j in j_n_list_array) + + μ = list(filter(lambda x: x[2] != 0.0 and not np.isnan(x).any(), mu)) + μ_x = list(map(lambda x: x[0], μ)) + μ_y = list(map(lambda x: x[1], μ)) + μ_alpha = list(map(lambda x: x[2], μ)) + true_μ_x = list(map(lambda x: x[0], true_μ)) + true_μ_y = list(map(lambda x: x[1], true_μ)) + true_μ_alpha = list(map(lambda x: x[2], true_μ)) + + if len(μ_alpha) == 0: + alpha_mi, alpha_ma, alpha_me = 0, 0, 0 + else: + alpha_mi, alpha_ma, alpha_me = min(μ_alpha), max(μ_alpha), median(μ_alpha) + + true_alpha_mi, true_alpha_ma = min(true_μ_alpha), max(true_μ_alpha) + alpha_base = min(true_alpha_mi, alpha_mi) + alpha_scale = max(true_alpha_ma, alpha_ma) - alpha_base + if alpha_scale == 0: + ms = lambda m: 6 + else: + ms = lambda m: int(1 + (m - alpha_base) / alpha_scale * 10) + + def plot_array( + name, t_idx, ax, u_array, u_min, u_max, colorbar=None, measure=False + ): + if u_min == u_max: + levels = [u_min, u_min + 1e-9] + else: + levels = np.linspace(u_min, u_max, quantisation) + try: + contour = ax.tricontourf(triang, u_array, levels=levels, cmap="viridis") + if colorbar: + colorbar.update_normal(contour) + except Exception as e: + print(e) + if measure: + for x, y, m in zip(μ_x, μ_y, μ_alpha): + ax.plot([x], [y], "ro", markersize=ms(m), label="Sources") + for x, y, m in zip(true_μ_x, true_μ_y, true_μ_alpha): + ax.plot([x], [y], "kx", markersize=ms(m), label="True sources") + if t_idx >= 0: + ax.set_title(f"%s; t = {t_idx:.1f}" % name) + else: + ax.set_title("%s" % name) + ax.set_aspect("equal") + + n = len(u_n_list_array) + # frames = list(map(lambda i: i * (n - 1) // (m - 1), range(0, m))) + frames = range(0, n) + for i, axu, axw, axj, t_idx in zip( + frames, + axus, + axws, + axjs, + times, # map(lambda i: i / (n - 1), frames), + ): + plot_array( + "u", t_idx, axu, u_n_list_array[i].real, u_min, u_max, colorbar_u, True + ) + plot_array( + "w", t_idx, axw, w_n_list_array[i].real, w_min, w_max, colorbar_w, True + ) + plot_array( + "j", t_idx, axj, j_n_list_array[i].real, j_min, j_max, colorbar_j, True + ) + + if show_true: + for i, ax, t_idx in zip( + frames, + ax_wpbar_plus, + map(lambda i: i / (n - 1), frames), + ): + plot_array( + "û", + t_idx, + ax, + true_u_n_list_array[i].real, + u_min, + u_max, + colorbar_u, + True, + ) + else: + plot_array("w̄ₚ", -1, ax_wpbar_plus[0], wp_bar_array.real, w_min, w_max) + plot_array("ω₀", -1, ax_wpbar_plus[1], ω0.real, w_min, w_max) + # plot_array("ω̂", -1, ax_wpbar_plus[2], true_ω.real, w_min, w_max) + + plt.suptitle( + "Convection-Diffusion, iteration %d; len(μ) = %d; μ_min = %f, μ_max = %f, μ_median = %f" + % (iter, len(μ), alpha_mi, alpha_ma, alpha_me), + fontsize=14, + ) + # plt.savefig("solution_evolution_%d.png" % iter, dpi=300) + + state = {"k": 0, "show_true": False} + + def on_key(event): + k0 = state["k"] + k = None + if event.key == "right" or event.key == " ": + k = (k0 + 1) % len(iter_files) + elif event.key == "left" or event.key == "backspace": + k = (k0 - 1) % len(iter_files) + elif event.key == "shift+right": + k = (k0 + 10) % len(iter_files) + elif event.key == "shift+left": + k = (k0 - 10) % len(iter_files) + elif event.key == "up": + k = (k0 + 100) % len(iter_files) + elif event.key == "down": + k = (k0 - 100) % len(iter_files) + elif event.key == "0": + k = 0 + elif event.key == "t": + state["show_true"] = not state["show_true"] + k = k0 + elif event.key == "q": + sys.exit() + if k is not None: + state["k"] = k + do_plot(iter_files[k], state["show_true"]) + fig.canvas.draw() + + do_plot(iter_files[0]) + + fig.canvas.mpl_connect("key_press_event", on_key) + + plt.show() + # # Time evolution + # fig, ax = plt.subplots(figsize=(10, 4)) + # times = np.arange(len(u_n_list)) * pde.dt + # max_vals = [np.max(u.x.array.real) for u in u_n_list] + + # ax.plot(times, max_vals, "r-o", linewidth=2, markersize=4) + # ax.set_xlabel("Time t") + # ax.set_ylabel("max|u|") + # ax.grid(True, alpha=0.3) + # ax.set_title("Solution Evolution") + # plt.tight_layout() + # plt.savefig("solution_time.png", dpi=150) + # plt.show() + + # print("Saved: solution_evolution.png + solution_time.png") + + +plot(sys.argv[1])
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/rustfmt.toml Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,3 @@ +overflow_delimited_expr = true +struct_lit_width = 80 +
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/compose.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,128 @@ +import numpy as np + + +class SumOfSeparableFunctions: + def __init__(self, fnlist): + self.fnlist = fnlist + + def apply(self, x): + val = 0.0 + for f_i, x_i in zip(self.fnlist, x): + val += f_i.apply(x_i) + return val + + def diff(self, x): + d = [] + for f_i, x_i in zip(self.fnlist, x): + d.append(f_i.diff(x_i)) + return d + + def apply_and_diff(self, x): + d = [] + val = 0.0 + for f_i, x_i in zip(self.fnlist, x): + (a, v) = f_i.apply_and_diff(x_i) + val += a + d.append(v) + return (val, d) + + def diff_lipschitz_factor(self): + res = 0 + for f_i in self.fnlist: + res = max(res, f_i.diff_lipschitz_factor()) + return res + + def diff_bound(self, xbound=None): + res = 0 + for f_i in self.fnlist: + res = max(res, f_i.diff_bound(xbound=xbound)) + return res + + +class ComposeFnWithOperator: + def __init__(self, f, op, xbound=None, xbound_pair=None): + self.f = f + self.op = op + self.xbound = xbound + self.xbound_pair = xbound_pair + + def apply(self, *args): + return self.f.apply(self.op.apply(*args)) + + def diff(self, *args): + # TODO: precalculations in apply should be used in diff_adjdir + w = self.op.apply(*args) + v = self.f.diff(w) + return self.op.diff_adjdir(v, *args, apply_result=w) + + def apply_and_diff(self, *args): + # TODO: precalculations in apply should be used in diff_adjdir + w = self.op.apply(*args) + (a, v) = self.f.apply_and_diff(w) + return (a, self.op.diff_adjdir(v, *args, apply_result=w)) + + def diff_lipschitz_factor(self): + # ‖∇A(x)^*∇F(A(x)) - ∇A(y)^*∇F(A(y))‖ + # = ‖[∇A(x)^*-∇A(y)^*]∇F(A(x)) - ∇A(y)^*[∇F(A(y))-∇F(A(x))]‖ + # ≤ L_{∇A(x)} M_{∇F} + M_{∇A(y)^*} L_{∇F}L_A. + if hasattr(self.op, "opnorm"): + # Linear operator + lda = 0.0 # This is zero, + mdf = 0.0 # hence this not needed. + else: + mdf = self.f.diff_bound(xbound=self.op.codomain_bound(xbound=self.xbound)) + lda = self.op.diff_adj_lipschitz_factor() + + ldf = self.f.diff_lipschitz_factor() + la = self.op.lipschitz_factor() + mda = self.op.diff_bound(xbound=self.xbound) + + return lda * mdf + mda * ldf * la + + def diff_lipschitz_factor_pair(self): + if self.op.hasattr("opnorm"): + # Linear operator + lda1, lda2 = 0.0, 0.0 # This is zero, + mdf = 0.0 # hence this not needed. + else: + lda1, lda2 = self.op.diff_adj_lipschitz_factor_pair() + mdf = self.f.diff_bound( + xbound=self.op.codomain_bound_pair(xbound=self.xbound_pair) + ) + + ldf = self.f.diff_lipschitz_factor() + la1, la2 = self.op.lipschitz_factor_pair() + mda = self.op.diff_bound_pair(xbound=self.xbound_pair) + + return lda1 * mdf + mda * ldf * la1, lda2 * mdf + mda * ldf * la2 + + +class InjectSecond: + def __init__(self, y): + self.y = y + + def apply(self, x): + return (x, self.y) + + def diff_adjdir(self, j, _x, apply_result=None): + return j[0] + + # This is not really a linear operator, but for our purposes affine behave essentially + # the same + def opnorm(self, *args): + return 1.0 + + def lipschitz_factor(self, *args): + return 1.0 + + def diff_adj_lipschitz_factor(self, *args): + return 0.0 + + def diff_bound(self, xbound=None): + return 1.0 + + def codomain_bound(self, xbound=None): + if xbound is None: + raise Exception("Linear operators have unbounded range") + else: + return xbound
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/convection_diffusion.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,894 @@ +import os +from dataclasses import dataclass +from posix import mkdir +from termios import B0 +from typing import Optional, Tuple, Union + +import dolfinx.geometry as geo +import numpy as np +import ufl +from dolfinx import fem, mesh +from dolfinx.fem.petsc import ( + apply_lifting, + assemble_matrix, + assemble_vector, + create_vector, + set_bc, +) +from mpi4py import MPI +from petsc4py import PETSc + +try: + import pointsource_pde.dolfinx_extras as dx + + from dolfinx_access import ( + cell_FunctionSpace_f64, + max_Function_f64_p2, + min_Function_f64_p2, + ) +except: + import src.dolfinx_extras as dx + + +@dataclass +class XBound: + # Parameter bounds for convection-diffusion operator estimates + mu_dual: float = 3.0 # ||μ||_M(Ω) bound + k_min: float = 0.1 # min diffusion > 0 (coercivity) + k_max: float = 2.0 # max diffusion + c_Linf: float = 0.5 # max(|c1|,|c2|) L∞ bound + diam: float = 1.41 # domain diameter (√2 for unit square) + T: float = 1.0 # final time + + def __add__(self, other: "XBound") -> "XBound": + # Conservative combination: min(k_min), max(c_Linf) for energy bounds""" + other = XBound( + **{ + k: getattr(other, k, getattr(self, k)) + for k in self.__dataclass_fields__ + } + ) + return XBound( + mu_dual=max(self.mu_dual, other.mu_dual), + k_min=min(self.k_min, other.k_min), + k_max=max(self.k_max, other.k_max), + c_Linf=max(self.c_Linf, other.c_Linf), + diam=max(self.diam, other.diam), + T=max(self.T, other.T), + ) + + +class ConvectionDiffusion: + def __init__( + self, + u0, + g, + # w0, + domain_size=[-2, 2, -2, 2], + nx=32, + ny=32, + degree=2, + t0=0.0, + T=1.0, + num_steps=50, + p=None, + Delta_t=None, + C_stab=1.0, + alpha=1.0, # \alpha = k_m/2, k convective coefficient, k>=km + beta=1.0, # \beta = \| c \|_\infty^2 /k_m + C_emb=1.0, + C_reg=1.0, + ): + self.p = p # L^p, p=2 + self.C_stab = C_stab # adjoint stability constant + self.alpha = alpha # coercivity lower bound for k + self.beta = beta # theoretical parameter + self.C_emb = C_emb # embedding constant + self.C_reg = C_reg + if Delta_t is None: + self.Delta_t = T / num_steps # Matches your N=50 intent + else: + self.Delta_t = Delta_t + + self.domain_size = domain_size + self.nx = nx + self.ny = ny + self.degree = degree + self.T = T + + self.num_steps = num_steps + self.dt = (T - t0) / num_steps + self.t0 = t0 + + self.save_for_plot = False + + domain = mesh.create_rectangle( + MPI.COMM_SELF, + [ + np.array([domain_size[0], domain_size[2]]), + np.array([domain_size[1], domain_size[3]]), + ], + [nx, ny], + mesh.CellType.triangle, + ) + self.domain = domain + # Compute the domain volume for the adjoint Lipschitz factor + self.domain_size = domain_size + # Domain diameter from your domain_size + self.diam = np.sqrt( + (domain_size[1] - domain_size[0]) ** 2 + + (domain_size[3] - domain_size[2]) ** 2 + ) + + V = fem.functionspace(domain, ("Lagrange", degree)) + self.V = V + + self.u0 = fem.Function(V) + self.u0.name = "u0" + self.u0.interpolate(u0) + + self.g = fem.Function(V) + self.g.name = "g" + self.g.interpolate(g) + self.bcs = [] # Initialize BCS as LIST + self.adjoint_bcs = [] # Initialize ADJOINT BCS as LIST (w = 0) + + # Identify boundary facets for Dirichlet BC location (e.g. left/right boundaries at x=domain_size[0] or x=domain_size[1]) + fdim = domain.topology.dim - 1 + # Γ₁ facets: TOP/BOTTOM boundaries (y = ymin, ymax) - STATE Dirichlet u = g + gamma1_facets = mesh.locate_entities_boundary( + domain, + fdim, + lambda x: ( + np.isclose(x[1], domain_size[2]) # x[1] = y = ymin + | np.isclose(x[1], domain_size[3]) + ), # x[1] = y = ymax + ) + # Γ₂ facets: LEFT/RIGHT boundaries (x = xmin, xmax) - ADJOINT Dirichlet w = 0 + gamma2_facets = mesh.locate_entities_boundary( + domain, + fdim, + lambda x: ( + np.isclose(x[0], domain_size[0]) # x[0] = x = xmin + | np.isclose(x[0], domain_size[1]) + ), # x[0] = x = xmax + ) + # STATE BC: u = g on Γ₁ (top/bottom) + dofs_gamma1 = fem.locate_dofs_topological(V, fdim, gamma1_facets) + # Create BC and ADD TO LIST + bc = fem.dirichletbc(self.g, dofs_gamma1) # V inferred from self.g! + self.bcs.append(bc) + + # ADJOINT BC: w = 0 on Γ₁ (same facets, homogenized) + zero = fem.Constant(domain, PETSc.ScalarType(0.0)) + dofs_gamma2 = fem.locate_dofs_topological(V, fdim, gamma2_facets) + bc_adjoint = fem.dirichletbc(zero, dofs_gamma2, V) + self.adjoint_bcs.append(bc_adjoint) + + self.fdim = fdim + self.boundary_facets = np.unique(np.concatenate((gamma1_facets, gamma2_facets))) + + ux = fem.Function(V) + wpx = fem.Function(V) + self.ux = ux + self.wpx = wpx + dim = V.mesh.topology.dim + self.expr2_form = fem.form(ufl.inner(ufl.grad(ux), ufl.grad(wpx)) * ufl.dx) + self.expr3_forms = tuple( + fem.form(wpx * ufl.grad(ux)[i] * ufl.dx) for i in range(dim) + ) + + # Only scalar case + self.k = fem.Constant(domain, PETSc.ScalarType(0.0)) # c[0] scalar + self.c0 = fem.Constant(domain, PETSc.ScalarType(0.0)) # c[0] scalar + self.c1 = fem.Constant(domain, PETSc.ScalarType(0.0)) # c[1] scalar + u_n = fem.Function(V) # New copy! + u_n.name = "u_n" + self.u_n = u_n + u, w, v = ufl.TrialFunction(V), ufl.TrialFunction(V), ufl.TestFunction(V) + dt = self.dt + # Forward bilinear form (backward Euler) + a = ( + u * v * ufl.dx + + dt * ufl.dot(self.k * ufl.grad(u), ufl.grad(v)) * ufl.dx + + dt * (self.c0 * u.dx(0) + self.c1 * u.dx(1)) * v * ufl.dx + ) + self.bilinear_form = fem.form(a) + L = u_n * v * ufl.dx + self.linear_form = fem.form(L) + + # Adjoint bilinear form + w_n = fem.Function(V) # New copy! + w_n.name = "w_n" + self.w_n = w_n + j_u = fem.Function(V) # New copy! + j_u.name = "j_u" + self.j_u = j_u + a2 = ( + w * v * ufl.dx + + dt + * ( + ufl.dot(self.k * ufl.grad(w), ufl.grad(v)) + - (self.c0 * w.dx(0) * v + self.c1 * w.dx(1) * v) + ) + * ufl.dx + ) + self.bilinear_form_w = fem.form(a2) + # Step n=N-1: L2 uses w_n = w^N=0 → solves w^{N-1}. (∂t)_primal^* + # Step n=N-2: L2 uses w_n = w^{N-1} → solves w^{N-2} + # Replace your L2 with 2 simple forms (same notation) + + L2 = w_n * v * ufl.dx + dt * j_u * v * ufl.dx # kown w_n (previous step) + self.linear_form_w = fem.form(L2) + + # Solving forward PDE + def apply(self, x): + μ, (k, c) = x + + domain = self.domain + V = self.V + dt = self.dt + bcs = self.bcs + + # initial condition + u_n = self.u_n + u_n.x.array[:] = self.u0.x.array + u_n.x.scatter_forward() + + # Linear form: only contains u_n * v (Dirichlet on Γ1 is handled via bc) + + # Assemble matrix WITH Dirichlet BCs applied + self.c0.value = c[0] + self.c1.value = c[1] + self.k.value = k + linear_form = self.linear_form + bilinear_form = self.bilinear_form + A = assemble_matrix(bilinear_form, bcs=bcs) + A.assemble() + + # Create reusable RHS vector (will be filled each timestep) + b = create_vector(linear_form) + b_ps = create_vector(linear_form) + + # Prepare solver once + solver = PETSc.KSP().create(domain.comm) + solver.setOperators(A) + solver.setType(PETSc.KSP.Type.PREONLY) + solver.getPC().setType(PETSc.PC.Type.LU) + + # Form point source contribution + mesh = V.mesh + element = V.element.basix_element + mesh_nodes = mesh.geometry.x + cmap = mesh.geometry.cmap + + b_ps.assemblyBegin() + + with b_ps.localForm() as loc_b: + loc_b.set(0) + + for point, alpha in μ.iter_padded(): + cell_id = cell_FunctionSpace_f64(V._cpp_object, point) + + if cell_id < 0: + print(f"Point {point} is outside the mesh.") + else: + cell_dofs = V.dofmap.cell_dofs(cell_id) + + geom_dofs = mesh.geometry.dofmap[cell_id] + ref = cmap.pull_back( + np.array([[point[0], point[1]]]), mesh_nodes[geom_dofs] + ) + phi = element.tabulate(0, ref) + + b_ps.setValuesLocal( + cell_dofs, + alpha * phi, + addv=PETSc.InsertMode.ADD_VALUES, + ) + + b_ps.assemblyEnd() + + t = self.t0 + num_steps = self.num_steps + bcs = self.bcs + + u_n_list = [] + for i in range(num_steps): + t += dt + + # b.assemblyBegin() + + # zero-out and assemble RHS (M * u_n part) + with b.localForm() as loc_b: + loc_b.set(0) + assemble_vector(b, linear_form) # b = (u^n, v) + + b.axpy(1.0, b_ps) + + # Apply homogeneous Dirichlet BC correction to RHS and finalize vector + apply_lifting(b, [bilinear_form], [bcs]) # [[bc1, bc2, ...]] + b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) + set_bc(b, bcs) # list [bc1, bc2, ...] + # b.assemble() + # + # b.assemblyEnd() + + # b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) + + # Solve linear problem A * uh = b + uh = fem.Function(V) + uh.name = "uh" + solver.solve(b, uh.x.petsc_vec) + uh.x.scatter_forward() + + # Update u_n for next step, store uh + u_n.x.array[:] = uh.x.array + u_n.x.scatter_forward() + u_n_list.append(uh) # Unique objects + + return u_n_list + + # Solving the adjoint problem + + def solve_adjoint_pde(self, j, x, u_n_list): + μ, (k, c) = x + + domain = self.domain + V = self.V + dt = self.dt + adjoint_bcs = self.adjoint_bcs + + num_steps = len(u_n_list) + + assert len(j) == num_steps, f"MISMATCH: j({len(j)}) != u_n_list({num_steps})" + # Terminal condition: w(T) = 0 + w_n = self.w_n + w_n.x.array[:] = 0.0 + w_n.x.scatter_forward() + + self.c0.value = c[0] + self.c1.value = c[1] + self.k.value = k + j_u = self.j_u + bilinear_form_w = self.bilinear_form_w + linear_form_w = self.linear_form_w + + A2 = assemble_matrix(bilinear_form_w, bcs=adjoint_bcs) # Fixed bcs + A2.assemble() + + # Create vector for RHS to be updated each step + b2 = create_vector(linear_form_w) + + solver = PETSc.KSP().create(domain.comm) + solver.setOperators(A2) + solver.setType(PETSc.KSP.Type.PREONLY) + solver.getPC().setType(PETSc.PC.Type.LU) + + t = self.t0 + num_steps * dt # Start at time T + + # N+1 adjoint values to match N PDE steps + w_n_list = [None] * (num_steps + 1) + whT = fem.Function(V) + whT.x.array[:] = 0.0 # w^N = 0 + whT.x.scatter_forward() + w_n_list[num_steps] = whT # Store at FINAL index + + # backward steps**: compute w^{N-1}, ..., w^0 + + for i in range(num_steps - 1, -1, -1): + t -= dt + + # Source term at current time step: j(t_i) + + j_u.x.array[:] = j[i].x.array + j_u.x.scatter_forward() + + # Update RHS + with b2.localForm() as loc_b2: + loc_b2.set(0) + assemble_vector(b2, linear_form_w) + + apply_lifting(b2, [bilinear_form_w], [adjoint_bcs]) # [forms] → [bcs] + b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) + set_bc(b2, adjoint_bcs) + + # set_bc(b2, adjoint_bcs) # list BC object + # b2.assemble() + + wh = fem.Function(V) + wh.name = "wh" + solver.solve(b2, wh.x.petsc_vec) + wh.x.scatter_forward() + + # Update w_n for next iteration and store + w_n.x.array[:] = wh.x.array + w_n.x.scatter_forward() + w_n_list[i] = wh + + return w_n_list + + # Compute the differential operator + def diff_adjdir(self, j, x, apply_result=None): + """ + Returns dual space elements: (dual_μ, (dual_k, (dual_c1, dual_c2))) + - scalar coeff → float sensitivity + - Function coeff → fem.Function sensitivity (Fréchet derivative) + """ + u_n_list = apply_result if apply_result is not None else self.apply(x) + w_n_list = self.solve_adjoint_pde(j, x, u_n_list) + dt = self.dt + V = self.V + dim = V.mesh.topology.dim + + # Extract coefficients and check types + _mu, (k, c) = x + # is_scalar_mu = isinstance(mu, (int, float, np.number)) + is_scalar_k = isinstance(k, (int, float, np.number)) + is_scalar_c = all(isinstance(ci, (int, float, np.number)) for ci in c) + + # 1. expr1: dual sensitivity w.r.t. μ = ∫₀^T w_p dt (always Function) + wp_bar = fem.Function(V) + wp_bar.x.array[:] = ( + 0.5 + * dt + * ( + w_n_list[0].x.array + + w_n_list[-1].x.array + + 2 * np.sum([w.x.array for w in w_n_list[1:-1]], axis=0) + ) + ) + + ux = self.ux + wpx = self.wpx + + # 2. expr2: dual sensitivity w.r.t. k + if is_scalar_k: + # k scalar → ⟨expr2, δk⟩ = δk ∫ ∇u·∇w_p → scalar + expr2 = 0.0 + + form = self.expr2_form + + for n, (u, wp) in enumerate(zip(u_n_list, w_n_list)): + ux.x.array[:] = u.x.array[:] + ux.x.scatter_forward() + wpx.x.array[:] = wp.x.array[:] + wpx.x.scatter_forward() + v2 = fem.assemble_scalar(form) + weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt + expr2 -= weight * v2 + + else: + raise Exception("Unimplemented - out of date") + + # k Function → ⟨expr2, δk⟩ = ∫ expr2 δk → expr2 = ∇u·∇w_p (pointwise) + expr2 = fem.Function(V) + expr2a = expr2.x.array + + for n, (u, wp) in enumerate(zip(u_n_list, w_n_list)): + form = ufl.inner(ufl.grad(u), ufl.grad(wp)) * ufl.dx + v2 = fem.assemble_scalar(fem.form(form)) + weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt + expr2a -= weight * v2 + + # 3. expr3: dual sensitivity w.r.t. c = (c1, c2) + if is_scalar_c: + # Scalar c → return tuple of 2 floats + + forms = self.expr3_forms + + def val(i): + expr3i = 0.0 + for n, (u, wp) in enumerate(zip(u_n_list, w_n_list)): + ux.x.array[:] = u.x.array[:] + ux.x.scatter_forward() + wpx.x.array[:] = wp.x.array[:] + wpx.x.scatter_forward() + v3 = fem.assemble_scalar(forms[i]) + weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt + expr3i -= weight * v3 + return expr3i + + expr3 = tuple(val(i) for i in range(dim)) + + else: + raise Exception("Unimplemented - out of date") + + # Function c → return tuple of 2 Functions (pointwise derivatives) + def val(i): + # expr3[i]. x.array[:] = 0 + expr3i = fem.Function(V) + expr3ia = expr3i.x.array[:] + for n, (u, wp) in enumerate(zip(u_n_list, w_n_list)): + form = wp * ufl.grad(u)[i] * ufl.dx + v3 = fem.assemble_scalar(fem.form(form)) + weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt + expr3ia -= weight * v3 + return expr3i + + expr3 = tuple(val(i) for i in range(dim)) + + if self.save_for_plot: + self.plot_wp_bar = wp_bar + self.plot_u_n_list = u_n_list + self.plot_w_n_list = w_n_list + self.plot_j_n_list = j + + return wp_bar, (expr2, expr3) + + def do_plot(self, iter, μ, frames, prefix): + n = self.num_steps + np.savez_compressed( + "%s/res_%d.npz" % (prefix, iter), + wp_bar_array=self.plot_wp_bar.x.array, + u_n_list_array=[self.plot_u_n_list[i].x.array for i in frames], + w_n_list_array=[self.plot_w_n_list[i].x.array for i in frames], + j_n_list_array=[self.plot_j_n_list[i].x.array for i in frames], + frames=frames, + times=list(map(lambda i: self.T * i / (n - 1), frames)), + mu=[[point[0], point[1], alpha] for point, alpha in μ.iter_padded()], + ) + + def _stability_prefactor( + self, *, C_stab=None, alpha=None, beta=None, T=None, C_emb=None, C_reg=None + ): + """ + Return the stability prefactor + P := (C_stab/alpha) * (1 + sqrt(T)*(1 + sqrt(beta)/sqrt(alpha))) + and the measure-to-functional constant L_e_mu := C_emb * C_reg * sqrt(T). + Return (P, L_e_mu): + """ + # self._stability_prefactor(alpha=0.5, T=2.0) + + # Read from arguments or instance attributes + C_stab = C_stab if C_stab is not None else self.C_stab + alpha = alpha if alpha is not None else self.alpha + beta = beta if beta is not None else self.beta + T = T if T is not None else self.T + C_emb = C_emb if C_emb is not None else self.C_emb + C_reg = C_reg if C_reg is not None else self.C_reg + + # Stability factor P + + P = (C_stab / alpha) * (1 + np.sqrt(T) * (1 + np.sqrt(beta) / np.sqrt(alpha))) + L_e_mu = C_emb * C_reg * np.sqrt(T) + + return P, L_e_mu + + # Solution operator Lipschitz factor + def lipschitz_factor(self): + """ + Conservative Lipschitz factor for the forward solution operator mapping + (k,c,mu) -> u. Returns a single scalar C_Lip such that + ||u2 - u1|| <= C_Lip * (||k2-k1|| + ||c2-c1|| + ||mu2-mu1||_* ) + (we use a conservative form combining the measure and (k,c) pieces). + """ + P, L_e_mu = self._stability_prefactor() + # Conservative combined choice: multiply P by max(1, L_e_mu) + C_Lip = P * max(1.0, L_e_mu) + return C_Lip + + # Adjoint solution operator Lipschitz factor + """ + def diff_adj_lipschitz_factor(self): + + Compute adjoint-solution Lipschitz factor A for 2D domain. + + For p = 2, A = 1. + For p > 2, A = |Omega|^(1/2 - 1/p), where |Omega| is the mesh area. + + p = self.p + omega_vol = self.domain_volume + return 1.0 if p == 2 else omega_vol ** (0.5 - 1.0 / p) + """ + + def diff_adj_lipschitz_factor(self): + """ + Always return 1 for p=2 + """ + return 1.0 + + def diff_bound(self, xbound=None): + """ + Differential bound for the solution operator derivative. + + Returns the tuple: (L_e_mu, L_e_k, L_e_c) + + where + L_e_k = ||∇u||_{L^2(Ω_T)} + L_e_c = ||∇u||_{L^2(Ω_T)} + L_e_mu = C_emb * C_reg * sqrt(T) + """ + + # If grad_u_norm is None, compute it from u0 + grad_u_norm = getattr(self, "grad_u_norm", None) + if grad_u_norm is None: + u = self.u0 + V = self.V + # Compute L2 norm of gradient of u0 over the domain + grad_u = ufl.grad(u) + grad_u_sq = ufl.inner(grad_u, grad_u) + dx = ufl.Measure("dx", domain=self.domain) + # Integral of |∇u|^2 over Ω + grad_u_norm = fem.assemble_scalar(fem.form(grad_u_sq * dx)) ** 0.5 + # Optionally store it for later reuse + self.grad_u_norm = grad_u_norm + + return grad_u_norm + + def diff_bound_pair(self, xbound=None, Delta_t=None): + if xbound is None: + return self.diff_bound(xbound=None) + else: + xb_combined = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0] + + k_min = getattr(xb_combined, "k_min", self.alpha) + c_Linf = getattr(xb_combined, "c_Linf", np.sqrt(self.beta * k_min)) + diam = getattr(xb_combined, "diam", self.diam) + T = getattr(xb_combined, "T", self.T) + + Pe = c_Linf * diam / k_min + gamma = c_Linf**2 / k_min + # adaptive time step size Delta_t (Delta_t = 0.01 in the test) + if Delta_t is None: + Delta_t = min(0.1, k_min / c_Linf**2, 0.01 * T) # CFL-like + + # Accumulate N = T/Delta_t local steps + N = int(np.ceil(T / Delta_t)) + + # Per-step factors + exp_local = np.exp(0.5 * gamma * Delta_t) + C_adj_step = np.sqrt(Delta_t * (exp_local + Delta_t / k_min)) + C_state_step = np.sqrt(Delta_t) * np.sqrt(1 + Pe**2) + + # Total accumulated bound (geometric sum <= N * max_step) + C_adj_PDE = N * C_adj_step + C_state = N * C_state_step + + # Use the embedding constant + C_mu_adj = self.C_emb * C_adj_PDE # ∫φ δμ + C_k_adj = C_state * C_adj_PDE # ∫∇u·∇φ δk + C_c1_adj = C_state * C_adj_PDE * c_Linf # ∫∂x u φ δc1 + C_c2_adj = C_state * C_adj_PDE * c_Linf # ∫∂y u φ δc2 + + return C_mu_adj, C_k_adj, C_c1_adj, C_c2_adj + + # ||S(x)||_Y→X ≤ C_state ||μ||_M(Ω) [time-independent μ] + def codomain_bound(self, xbound=None): + if xbound is None: + return 1.0 + + # Extract parameters + if hasattr(xbound, "k_min"): + xb = xbound + else: + if hasattr(xbound, "__len__") and len(xbound) > 0: + xb = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0] + else: + xb = xbound + k_min = getattr(xb, "k_min", self.alpha) + c_Linf = getattr(xb, "c_Linf", np.sqrt(self.beta * k_min)) + mu_dual = getattr(xb, "mu_dual", 3.0) + diam = getattr(xb, "diam", self.diam) + T = getattr(xb, "T", self.T) + + Pe = c_Linf * diam / k_min + if not hasattr(self, "Delta_t"): + Delta_t = min(0.1, k_min / c_Linf**2, 0.01 * T) + else: + Delta_t = self.Delta_t + + N = int(np.ceil(T / Delta_t)) + gamma = c_Linf**2 / k_min + exp_local = np.exp(0.25 * gamma * Delta_t) + + C_state_step = np.sqrt(Delta_t) * np.sqrt(1 + 0.5 * Pe**2) * exp_local + C_state = N * C_state_step + + return C_state * mu_dual + + # ||S(x1)-S(x2)|| ≤ Cμ||Δμ|| + Ck||Δk|| + Cc1||Δc1|| + Cc2||Δc2|| + + def codomain_bound_pair(self, xbound=None): + if xbound is None: + return self.codomain_bound(xbound=None) + else: + xb = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0] + C_base = self.codomain_bound(xb) + # Component + C_mu = C_base + C_k = C_base * 2.0 + C_c1 = C_base * 1.5 + C_c2 = C_base * 1.5 + + return (C_mu, C_k, C_c1, C_c2) + + """ + def codomain_bound(self, xbound=None): + return 1.0 + + def codomain_bound_pair(self, xbound=None): + # TODO: implement properly + if xbound is None: + return self.codomain_bound(xbound=None) + else: + return self.codomain_bound(xbound=xbound[0] + xbound[1]) + """ + + def diff_bound3(self): + """ + Returns full operator differential bounds: + + L_e_mu = C_emb * C_reg * sqrt(T) + L_e_k = ||∇u||_{L^2(Ω)} + L_e_c = ||∇u||_{L^2(Ω)} + + """ + + # First get L_e_k = L_e_c + grad_u_norm = self.diff_bound() + + # Compute L_e_mu from stability prefactor + _P, L_e_mu = self._stability_prefactor() + + L_e_k = grad_u_norm + L_e_c = grad_u_norm + + return L_e_mu, L_e_k, L_e_c + + # Solution operator Lipschitz factor separately wrt. μ and (k, c) + def lipschitz_factor_pair(self): + """ + Compute the forward solution operator Lipschitz factors + separately with respect to the parameters (k, c) and μ. + + Returns: + tuple: (L_mu, L_kc) + L_mu - Lipschitz factor w.r.t. μ + L_kc - Lipschitz factor w.r.t. (k, c) + """ + # Get differential bounds for the solution operator + L_e_mu, L_e_k, L_e_c = self.diff_bound3() + + # Lipschitz factor w.r.t μ + L_mu = L_e_mu + + # Lipschitz factor w.r.t (k, c) + # Conservative choice: max of L_e_k and L_e_c + L_kc = max(L_e_k, L_e_c) + + return L_mu, L_kc + + # Adjoint solution operator Lipschitz factor separately wrt. μ and (k, c) + def diff_adj_lipschitz_factor_pair(self): + A = self.diff_adj_lipschitz_factor() + + # No separate analytic dependence, so return same A for both parts + A_mu = A + A_kc = A + + return A_mu, A_kc + + +def own(u): + # """ + # Return a version of u that is guaranteed to be uniquely owned. + # If u has other Python references, return a deep copy. + # """ + # # sys.getrefcount(u) includes the temporary reference inside getrefcount, + # # so '2' means exactly one external reference. + # if sys.getrefcount(u) <= 2: + # return u # safe: no other references + + # deep copy into a new Function + u_new = fem.Function(u.function_space) + # Does not work + # u_new.x.petsc_vec.copy(u.x.petsc_vec) + # u_new.x.scatter_forward() + np.copyto(u_new.x.array, u.x.array) + return u_new + + +def nn2(x): + return None if isinstance(x, float) else dx.norm2_squared(x) + + +class PlotFactory: + def __init__(self, pde, n_plots=5): + self.pde = pde + self.n_plots = n_plots + + m = n_plots + n = pde.num_steps + pde.save_for_plot = True + self.plot_frames = list(map(lambda i: i * (n - 1) // (m - 1), range(0, m))) + self.pde = pde + + def produce(self, prefix): + return Plotter(self.pde, self.n_plots, prefix) + + +class Plotter: + def __init__(self, pde, n_plots, prefix): + m = n_plots + n = pde.num_steps + pde.save_for_plot = True + self.plot_frames = list(map(lambda i: i * (n - 1) // (m - 1), range(0, m))) + self.pde = pde + os.makedirs(prefix, exist_ok=True) + self.prefix = prefix + + def plot(self, iter, μ): + self.pde.do_plot(iter, μ, self.plot_frames, self.prefix) + + +class QuadraticRegularisation: + """ + Quadratic regularisation functional for the auxiliary variables + """ + + def _own(self, x): + return x if isinstance(x, float) else own(x) + + def __init__(self, α, base=None): + self.α = α + if base is not None: + (k, (c1, c2)) = base + self.base = (self._own(k), (self._own(c1), self._own(c2))) + self.base_norm_squared = (nn2(k), (nn2(c1), nn2(c2))) + else: + self.base = None + + def _apply1(self, x, b, n): + if isinstance(x, float): + d = x if b is None else x - b + return d * d / 2.0 + else: + if b is not None: + return (dx.norm2_squared(x) + n) / 2.0 + dx.dot(x, b) + return dx.norm2_squared(x) / 2.0 + + def _get_base(self): + return (None, (None, None)) if self.base is None else self.base + + def apply(self, x): + """ + Apply the quadratic regularisation fucntional to `x`. + """ + (k, (c1, c2)) = x + (kb, (c1b, c2b)) = self._get_base() + (nkb, (nc1b, nc2b)) = self.base_norm_squared + + return self.α * ( + self._apply1(k, kb, nkb) + + self._apply1(c1, c1b, nc1b) + + self._apply1(c2, c2b, nc2b) + ) + + def _prox1(self, τ, x, b): + γ = self.α * τ + β = 1.0 / (1.0 + γ) + if isinstance(x, float): + return β * (x if b is None else x + b * γ) + else: + # x = own(x) + vx = x.x.array + if b is not None: + vx += b.x.array * γ + vx *= β + return x + + def prox(self, τ, x): + """ + Apply the proximal map of the quadratic regularisation fucntional to `x`. + WARNING: This function is unsafe. It may modify `x´. + That is ok for our purposes, as Rust, being a safe language, already needs to pass a + copied or owned instance to Python. + """ + k, (c1, c2) = x + (kb, (c1b, c2b)) = self._get_base() + + return ( + self._prox1(τ, k, kb), + (self._prox1(τ, c1, c1b), self._prox1(τ, c2, c2b)), + )
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/dolfinx_access.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,209 @@ +use alg_tools::error::DynResult; +use alg_tools::fe_model::{ + base::RealLocalModel, + p2_local_model::{P2LocalModel, Simplex}, +}; +use alg_tools::loc::Loc; +use anyhow::bail; +use pyo3::ffi::PyImport_AppendInittab; +use std::ffi::CString; +use std::mem::MaybeUninit; + +// These are required for the linking to the sys crates to actually happen. +#[allow(unused_imports)] +mod dummy_import { + use dolfinx_sys; + use nanobind_sys; +} + +mod function; +pub use function::DolfinxPyFunction_f64; + +#[allow(dead_code)] +#[cxx::bridge(namespace = "dolfinx_access")] +pub mod ffi { + #[derive(Copy, Clone, Debug)] + struct CoordValuePair { + x: [f64; 2], + v: f64, + } + + #[derive(Copy, Clone, Debug)] + struct FunctionInfo { + domain_dim: u32, + codomain_dim: u32, + order: u32, + triangular_mesh: bool, + } + + unsafe extern "C++" { + include!("pointsource_pde/include/dolfinx_access/nanobind_helpers.h"); + include!("pointsource_pde/include/dolfinx_access/function.h"); + include!("pointsource_pde/include/dolfinx_access/minmax_p2.h"); + + type Function_f64; + type PyObject; + + /// Find the cell containing `x` (in 1D, 2D, or 3D, following Fenics weirdness + /// of hard-coding vectors to 3D). + pub unsafe fn cell_Function_f64(f: *const Function_f64, x: &[f64; 3]) -> i32; + + /// Evaluate `f` at `x`. + pub unsafe fn eval_Function_f64(f: *const Function_f64, x: &[f64; 3]) -> f64; + + /// Evaluates `f` at `x` (in 1D, 2D, or 3D, following Fenics weirdness + /// of hard-coding vectors to 3D), assuming `x` to belong to `cell`. + pub unsafe fn eval_Function_f64_cell( + f: *const Function_f64, + x: &[f64; 3], + cell: i32, + ) -> f64; + + /// Evaluate the function at 6 coordinates known to belong to given cell + pub unsafe fn eval_Function_f64_cell_6( + f: *const Function_f64, + all_coords: &[f64; 18], // 6*3 + cell: i32, + ) -> [f64; 6]; + + pub unsafe fn info_Function_f64(f: *const Function_f64) -> FunctionInfo; + + pub unsafe fn min_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>; + pub unsafe fn max_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>; + pub unsafe fn minmax_Function_f64_p2( + f: *const Function_f64, + max: bool, + ) -> Result<CoordValuePair>; + + pub unsafe fn check_Function_f64(o: *const PyObject) -> bool; + pub unsafe fn cast_Function_f64(o: *const PyObject) -> Result<*const Function_f64>; + pub unsafe fn cast_mut_Function_f64(o: *mut PyObject) -> Result<*mut Function_f64>; + pub unsafe fn cell_coords_Function_f64_triangle( + f: *const Function_f64, + cell: i32, + ) -> [f64; 9]; + + /// Returns the local weights + pub unsafe fn data_Function_f64<'a>(f: *const Function_f64) -> &'a [f64]; + + /// Returns the local mutable weights + pub unsafe fn data_mut_Function_f64<'a>(f: *mut Function_f64) -> &'a mut [f64]; + + /// Check the compatibility of function spaces of two functions. + pub unsafe fn check_compat_Function_f64( + f: *const Function_f64, + g: *const Function_f64, + ) -> Result<()>; + + /// Create a new similar function + pub unsafe fn similar_Function_f64(f: *const Function_f64) -> *mut Function_f64; + + /// Clone the function + pub unsafe fn clone_Function_f64(f: *const Function_f64) -> *mut Function_f64; + + pub unsafe fn wrap_Function_f64(f: *mut Function_f64) -> *mut PyObject; + + pub unsafe fn drop_Function_f64(f: *mut Function_f64); + } + + extern "Rust" { + unsafe fn minmax_dolfinx_p2_cell( + f: *const Function_f64, + simplex_coords: &[f64; 9], + cell: i32, + max: bool, + ) -> CoordValuePair; + } +} + +use ffi::CoordValuePair; + +/// Convert an array of 2D oordinates obtained from Fenics, which hard-codes weird 3D coordinates. +#[inline] +fn coords_from_fenics(&[x1, x2, _, y1, y2, _, z1, z2, _]: &[f64; 9]) -> [Loc<2, f64>; 3] { + [Loc([x1, x2]), Loc([y1, y2]), Loc([z1, z2])] +} +// fn coords_from_fenics<const N: usize>(coords: &[[f64; 3]; N]) -> [Loc<F, 2>; N] { +// coords.map(|[x1, x2, _]| Loc([x1, x2])) +// } + +/// Convert an array of 2D oordinates to Fenics, which hard-codes weird 3D coordinates. +#[inline] +fn coords_to_fenics([Loc([x1, x2]), Loc([y1, y2]), Loc([z1, z2])]: [Loc<2, f64>; 3]) -> [f64; 9] { + [x1, x2, 0.0, y1, y2, 0.0, z1, z2, 0.0] +} +// fn coords_to_fenics<const N: usize>(coords: [Loc<F, 2>; N]) -> [[f64; 3]; N] { +// coords.map(|Loc([x1, x2])| [x1, x2, 0.0]) +// } + +/// Given a dolphinx `Function` `f`, as well as 6 coordinates on a simplex in its mesh +/// (the corners and midpoints), construct a new internal P2 model on `cell`. +/// If `max` is true, the model is negated compared to `f`. +unsafe fn model_dolfinx_p2_cell( + f: *const ffi::Function_f64, + simplex_coords: &[f64; 9], + cell: i32, + max: bool, +) -> (P2LocalModel<f64, 2, 3>, Simplex<f64, 2, 3>) { + // Form simplex + let simplex = Simplex(coords_from_fenics(simplex_coords)); + // Form an array of all coordinates + let midpoints = simplex.midpoints(); + let mut all_coords_uninit: [MaybeUninit<f64>; 6 * 3] = [const { MaybeUninit::uninit() }; 6 * 3]; + all_coords_uninit + .iter_mut() + .zip( + simplex_coords + .iter() + .copied() + .chain(coords_to_fenics(midpoints).into_iter()), + ) + .for_each(|(a, b)| { + a.write(b); + }); + let all_coords = MaybeUninit::array_assume_init(all_coords_uninit); + // Evaluate the function, and invert the signs of maximising + let mut values = ffi::eval_Function_f64_cell_6(f, &all_coords, cell); + if max { + values.iter_mut().for_each(|v| *v = -*v); + } + // Form local model + let [ref n0, ref n1, ref n2] = simplex.0; + let [ref n01, ref n12, ref n20] = midpoints; + let nodes = [*n0, *n1, *n2, *n01, *n12, *n20]; + let model = P2LocalModel::<f64, 2, 3>::new(&nodes, &values); + + return (model, simplex); +} + +/// This is a helper routine for `minmax_dolfinx_p2`. It reconstructs a P2 +/// model on the cell by evaluating f at 6 points (the simplex corners given +/// in `simplex_doords`, as well as edge midpoints). Then it minimises this +/// model. The result should be accurate, given that `f` is assumed to have a +/// P2 model on the simplex. +pub unsafe fn minmax_dolfinx_p2_cell( + f: *const ffi::Function_f64, + simplex_coords: &[f64; 9], + cell: i32, + max: bool, +) -> ffi::CoordValuePair { + let (model, simplex) = model_dolfinx_p2_cell(f, simplex_coords, cell, max); + let (x, v) = model.minimise(&simplex); + CoordValuePair { x: x.into(), v } +} + +pub const NAME: &str = "dolfinx_access"; + +unsafe extern "C" { + pub fn PyInit__dolfinx_access() -> *mut pyo3::ffi::PyObject; +} + +/// Register the module in Python +pub fn register_python_ffi() -> DynResult<()> { + let cname = CString::new(NAME).unwrap(); + // We need to use into_raw() to transfer ownership; otherwise this stops working. + if unsafe { PyImport_AppendInittab(cname.into_raw(), Some(PyInit__dolfinx_access)) } != 0 { + bail!("Failed to add embedded nanobind dolfinx_access module"); + } + Ok(()) +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/dolfinx_access/function.cc Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,167 @@ +#include <array> +#include <cmath> +#include <span> + +#include <basix/finite-element.h> +#include <dolfinx/fem/Function.h> +#include <nanobind/nanobind.h> + +#include "dolfinx_access/function.h" +// #include "mpi_proto.h" +#include "dolfinx/geometry/BoundingBoxTree.h" +#include "pointsource_pde/src/dolfinx_access.rs.h" +#include "rust/cxx.h" + +using namespace dolfinx::fem; +using namespace dolfinx::geometry; +using namespace basix::element; +using namespace dolfinx::graph; +using namespace dolfinx::la; +namespace cell = basix::cell; + +namespace dolfinx_access { + extern void drop_Function_f64(Function_f64* f) { + delete f; + } + + FunctionInfo info_Function_f64(Function_f64 const* f) { + auto fp = f->function_space(); + auto el = fp->element()->basix_element(); + return FunctionInfo{ + .domain_dim = (uint32_t)fp->mesh()->geometry().dim(), + .codomain_dim = (uint32_t)fp->value_size(), + .order = (uint32_t)fp->element()->basix_element().degree(), + .triangular_mesh = el.cell_type() == cell::type::triangle, + }; + } + + void check_compat_Function_f64(Function_f64 const* f, Function_f64 const* g) { + if (f->function_space() != g->function_space()) { + throw "Incompatible function spaces"; + } + } + + std::map<std::weak_ptr<const mesh::Mesh<double>>, std::shared_ptr<BoundingBoxTree<double>>, + std::owner_less<std::weak_ptr<const mesh::Mesh<double>>>> + bb_tree_cache; + + /// Returns the cell containg x + int32_t cell_Mesh_f64(std::shared_ptr<const mesh::Mesh<double>> mesh, + const std::array<double, 3>& x) { + + std::weak_ptr<const mesh::Mesh<double>> mesh_weak = mesh; + std::shared_ptr<BoundingBoxTree<double>> bb_tree = bb_tree_cache[mesh_weak]; + if (bb_tree == nullptr) { + auto topo = mesh->topology(); + // auto dim = mesh->geometry().dim(); + auto dim = topo->dim(); + + // This fails as it inorrectly calls mesh.topology_mutable(), so need to + // copy-paste the implementation below, with the appropriate fix. + //*bb_tree = new BoundingBoxTree(*mesh, dim); + // And this also doesn't work due to the privacy of `range`. + //*bb_tree = new BoundingBoxTree(*mesh, dim, BoundingBoxTree::range(mesh->topology(), + // dim + auto index_map = topo->index_map(dim); + int local_cells = index_map->size_local(); + // We don't really do MPI; pointsource_algs is not designed for it. + assert(local_cells == index_map->size_global()); + std::vector<std::int32_t> r(local_cells); + std::iota(r.begin(), r.end(), 0); + bb_tree = std::make_shared<BoundingBoxTree<double>>(BoundingBoxTree(*mesh, dim, r)); + bb_tree_cache[mesh_weak] = bb_tree; + } + // Find colliding bounding boxes. This is an AdjancencyTree. + // + // C++ and Dolfinx are just vomit-inducing 🤮. We need some weird const double span here, + // but then just standard x in compute_first_collding_cell below. + auto colliding_bbs_adj = compute_collisions(*bb_tree.get(), std::span<const double>(x)); + // Since we compute the collisions for just one node, the `array` of all links + // of the adjancency list, will corresponding exactly the ccells whose bounding boxes + // collide with the point. + auto colliding_bbs = colliding_bbs_adj.array(); + // Within the colliding bounding boxes, find the first cell that actually collides. + int32_t local_cell = compute_first_colliding_cell(*mesh, colliding_bbs, x, 1e-9); + // local_cell may be -1 if it was not found by the local process, so combine + // results with MPI_MAX. + int32_t global_cell = local_cell; + MPI_Allreduce(&local_cell, &global_cell, 1, MPI_INT32_T, MPI_MAX, mesh->comm()); + + return global_cell; + } + + int32_t cell_FunctionSpace_f64(FunctionSpace<double> const* v, + const std::array<double, 3>& x) { + return cell_Mesh_f64(v->mesh(), x); + } + + int32_t cell_Function_f64(Function_f64 const* f, const std::array<double, 3>& x) { + return cell_Mesh_f64(f->function_space()->mesh(), x); + } + + std::array<double, 3 * 3> cell_coords_Function_f64_triangle(Function_f64 const* f, + int32_t cell) { + auto geom = f->function_space()->mesh()->geometry(); + auto gx = geom.x(); + auto geom_dofmap = geom.dofmap(); + assert(geom_dofmap.extent(1) == 3); + + auto j0 = geom_dofmap(cell, 0); + auto j1 = geom_dofmap(cell, 1); + auto j2 = geom_dofmap(cell, 2); + + // Construct list of coordinates. Due to f->eval, we construct this as a single list + // with all the Fenics weirdness of hard-coded 3D. + return {gx[3 * j0], gx[3 * j0 + 1], 0, /* */ + gx[3 * j1], gx[3 * j1 + 1], 0, /* */ + gx[3 * j2], gx[3 * j2 + 1], 0}; + } + + // We assume `f` to be real-valued. + double eval_Function_f64_cell(Function_f64 const* f, const std::array<double, 3>& x, + int32_t cell) { + if (cell < 0) { + return 0.0; + } else { + std::array<double, 1> res; + f->eval(x, {1, 3}, {{cell}}, std::span(res), {1, 1}); + return res[0]; + } + } + + double eval_Function_f64(Function_f64 const* f, const std::array<double, 3>& x) { + return eval_Function_f64_cell(f, x, cell_Function_f64(f, x)); + } + + std::array<double, 6> eval_Function_f64_cell_6(Function_f64 const* f, + const std::array<double, 3 * 6>& x, + int32_t cell) { + std::array<double, 6> res; + auto i = cell; + f->eval(x, {6, 3}, {{i, i, i, i, i, i}}, std::span(res), {6, 1}); + return res; + } + + rust::Slice<const double> data_Function_f64(Function_f64 const* f) { + auto span = f->x()->array(); + return rust::Slice(span.data(), span.size()); + } + + rust::Slice<double> data_mut_Function_f64(Function_f64* f) { + auto span = f->x()->mutable_array(); + return rust::Slice(span.data(), span.size()); + } + + Function_f64* similar_Function_f64(Function_f64 const* f) { + return new Function_f64(f->function_space()); + } + + Function_f64* clone_Function_f64(Function_f64 const* f) { + return new Function_f64(f->function_space(), std::make_shared<Vector<double>>(*f->x())); + } + + // PyObject* py_similar_Function_f64(Function_f64 const* f) { + // return nanobind::cast(Function_f64(f->function_space())).release().ptr(); + // } + +} // namespace dolfinx_access
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/dolfinx_access/function.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,812 @@ +/*! +Python and C++ wrapper for Dolfinx Function<f64> +*/ + +use super::ffi; +use super::ffi::Function_f64; +use alg_tools::bounds::{Bounds, GlobalAnalysis, MinMaxMapping}; +use alg_tools::euclidean::Euclidean; +use alg_tools::fe_model::base::LocalModel; +use alg_tools::instance::Instance; +use alg_tools::loc::Loc; +use alg_tools::mapping::{BasicDecomposition, DifferentiableImpl, Mapping, Space}; +use alg_tools::norms::{Dist, HasDual, Norm, Normed, L2}; +use itertools::izip; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyFunction}; +use std::ffi::CStr; + +/// A helper structure of dealing with dolfinx functions. +/// `N` is the domain dimension, `O` the order, and `D` is the codomain dimension. +#[allow(non_camel_case_types)] +pub struct DolfinxPyFunction_f64<'py, const N: u32, const O: u32 = 1, const D: u32 = 1> { + #[allow(dead_code)] + /// Python object. + /// + /// This is crucial for maintaining [`cxx`] live, when we obtained the object from Python. + /// + /// If we needed to copy-on-write, or initialise without obtaining and object from Python, + /// this will become `None`. + pub u: Option<Bound<'py, PyAny>>, + /// We need to maintain a reference to the FunctionSpace Python wrapper, as it's difficult + /// to recreate if we loose `py`. Fenics is a one-way hell full of pointless wrappers upon + /// wrappers upon wrappers,very difficult to create having with just the unwrapped object. + function_space: Bound<'py, PyAny>, + /// Do we have the only reference? + /// + /// This is check on construction only, when we have the GIL. + /// If we have the only reference at that time, this will not change during the lifetime. + /// We use this to implement copy-on-write optimisations. + owned: bool, + /// C++ object pointer + cxx: *mut Function_f64, +} + +struct Helpers { + dot: Py<PyFunction>, + norm2_squared: Py<PyFunction>, + dist2_squared: Py<PyFunction>, +} + +use std::sync::OnceLock; +static HELPERS: OnceLock<Helpers> = OnceLock::new(); + +fn get_helpers(py: Python<'_>) -> PyResult<&Helpers> { + HELPERS.get_or_try_init(|| { + let extras = PyModule::import(py, "pointsource_pde.dolfinx_extras")?; + let g = |s| -> PyResult<Py<PyFunction>> { Ok(extras.getattr(s)?.cast_into()?.unbind()) }; + Ok(Helpers { + dot: g("dot")?, + norm2_squared: g("norm2_squared")?, + dist2_squared: g("dist2_squared")?, + }) + }) +} + +impl<'py, const N: u32, const O: u32, const D: u32> Drop for DolfinxPyFunction_f64<'py, N, O, D> { + fn drop(&mut self) { + // If Python is not managing our C++ object, we need to free it + if let None = self.u { + if self.cxx != std::ptr::null_mut() { + unsafe { ffi::drop_Function_f64(self.cxx) }; + } + } + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> Clone for DolfinxPyFunction_f64<'py, N, O, D> { + fn clone(&self) -> Self { + DolfinxPyFunction_f64 { + u: None, // It's a different new function, so no Python instance + function_space: self.function_space.clone(), + owned: self.owned, + cxx: unsafe { ffi::clone_Function_f64(self.cxx) }, + } + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> Space for DolfinxPyFunction_f64<'py, N, O, D> { + type Principal = Self; + type Decomp = BasicDecomposition; +} + +impl<'py, const N: u32, const O: u32, const D: u32> DolfinxPyFunction_f64<'py, N, O, D> { + /// Create a new instance, assuming `cxx` has been correctly extracted from `u` + /// and checked to match `N`, `O`, `D`, and mesh restrictions. + pub(crate) fn new_prechecked(u: Bound<'py, PyAny>, cxx: *mut Function_f64) -> PyResult<Self> { + let function_space = u.getattr("_V")?; + let owned = u.get_refcnt() <= 1; + // TODO: check dimensionality + Ok(DolfinxPyFunction_f64 { u: Some(u), function_space, owned, cxx }) + } + + /// Ensures that `self.u` exists. + /// Panics if there's a failure to create it if needed. + pub(crate) fn _make_py(&mut self, py: Python<'py>) { + if let None = self.u { + // We need to do a lot of work here, due to Fenics one-wayness. + + // This takes ownership of our pointer, so if we fail to create the Python + // object, we have a lot of restoring to do. + let cpp_object = unsafe { + Py::<PyAny>::from_owned_ptr( + py, + ffi::wrap_Function_f64(self.cxx) as *mut pyo3::ffi::PyObject, + ) + }; + let locals = PyDict::new(py); + match locals.set_item("cpp_object", cpp_object).and_then(|()| { + locals.set_item("V", self.function_space.as_borrowed())?; + // TODO: precompile + py.run(CREATE_FUNCTION_HACK, None, Some(&locals))?; + locals.get_item("wrapper") + }) { + Err(e) => { + // We should restore to recover + panic!("Failure to pass object to Python: {:?}", e); + } + Ok(maybe_p) => { + self.u = Some(maybe_p.unwrap()); + } + } + } + } + + /// Returns a (possibly new) Python presentation, without consuming current object. + /// + /// Please use this functionality through [`pyo3::types::IntoPyObject`]. + pub(crate) fn _into_py_ref<'a>( + &'a mut self, + py: Python<'py>, + ) -> PyResult<pyo3::Borrowed<'a, 'py, PyAny>> { + self._make_py(py); + self.owned = false; + match self.u { + Some(ref p) => Ok(p.as_borrowed()), + None => unreachable!("This cannot happen"), + } + } + + /// Returns a (possibly new) Python presentation, consuming current object. + /// + /// Please use this functionality through [`pyo3::types::IntoPyObject`]. + pub(crate) fn _into_py(mut self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, PyAny>> { + self._make_py(py); + self.cxx = std::ptr::null_mut(); + self.owned = false; + match std::mem::replace(&mut self.u, None) { + Some(p) => Ok(p), + None => unreachable!("This cannot happen"), + } + } + + /// Create a new FEM function on the same function space, with unintialised data. + pub fn similar(&self) -> Self { + DolfinxPyFunction_f64 { + u: None, + function_space: self.function_space.clone(), + owned: true, + cxx: unsafe { ffi::similar_Function_f64(self.cxx) }, + } + } + + /// Create a new FEM function on the same function space, with zero data. + pub fn similar_zeros(&self) -> Self { + let new = self.similar(); + for x in unsafe { ffi::data_mut_Function_f64(new.cxx) } { + *x = 0.0; + } + new + } +} + +pub(crate) static CREATE_FUNCTION_HACK: &CStr = c_str!( + "\ +from dolfinx.fem import Function +from dolfinx.la import Vector +x = Vector(cpp_object.x) +# Fenics is very one-way. It does not provide a proper way to create the wrapper, +# directly from the unwrapped C++ +wrapper = Function.__new__(Function, V, x, \"?\", x.array.dtype) +# This all just repeats Function.__init__ without creating a new ._cpp_object. +super(Function, wrapper).__init__(V.ufl_function_space()) +wrapper._cpp_object = cpp_object; +wrapper._V = V +wrapper._x = x +wrapper.name = \"?\" +" +); + +trait ExpandCoordDolphinx { + fn to_dolphinx(self) -> [f64; 3]; + #[allow(dead_code)] + fn from_dolphinx(val: [f64; 3]) -> Self; +} + +impl ExpandCoordDolphinx for Loc<1, f64> { + #[inline] + fn to_dolphinx(self) -> [f64; 3] { + let Loc([x1]) = self; + [x1, 0.0, 0.0] + } + + #[inline] + fn from_dolphinx([x1, _, _]: [f64; 3]) -> Self { + [x1].into() + } +} + +impl ExpandCoordDolphinx for Loc<2, f64> { + #[inline] + fn to_dolphinx(self) -> [f64; 3] { + let Loc([x1, x2]) = self; + [x1, x2, 0.0] + } + + #[inline] + fn from_dolphinx([x1, x2, _]: [f64; 3]) -> Self { + [x1, x2].into() + } +} + +impl ExpandCoordDolphinx for Loc<3, f64> { + #[inline] + fn to_dolphinx(self) -> [f64; 3] { + let Loc([x1, x2, x3]) = self; + [x1, x2, x3] + } + + #[inline] + fn from_dolphinx([x1, x2, x3]: [f64; 3]) -> Self { + [x1, x2, x3].into() + } +} + +macro_rules! impl_diff { + ($n:literal ; $($o:literal)+) => { $( + impl<'py> DifferentiableImpl<Loc<$n, f64>> for DolfinxPyFunction_f64<'py, $n, $o, 1> { + type Derivative = Loc<$n, f64>; + + fn differential_impl<I: Instance<Loc<$n, f64>>>(&self, x: I) -> Self::Derivative { + let x_own = x.own(); + let f = self.cxx; + let expand_x = x_own.to_dolphinx(); + // Find cell containing `x` + let cell = unsafe { ffi::cell_Function_f64(f, &expand_x) }; + if cell < 0 { + Loc::ORIGIN + } else { + let (model, _) = unsafe { + // Find the coordinates of the corresponding triangle. We guarantee a + // triangular mesh in the constructor of `DolfinxPyFunction_f64`. + let simplex_coords = ffi::cell_coords_Function_f64_triangle(f, cell); + // Construct a new P2 model on that cell, as this is easier than dealing + // with the awful documentation of Dolphinx and C++, and probably having + // to reimplemen our P2 stuff anyway, to do any optimisation with the basix + // basis functions. + super::model_dolfinx_p2_cell(f, &simplex_coords, cell, false) + }; + model.differential(&x_own) + } + } + } + )+ }; +} + +macro_rules! impl_mapping { + ($($n:literal)+) => { $( + impl<'py, const O: u32> Mapping<Loc<$n, f64>> for DolfinxPyFunction_f64<'py, $n, O, 1> { + type Codomain = f64; + + fn apply<I: Instance<Loc<$n, f64>>>(&self, x: I) -> f64 { + let expand_x = x.own().to_dolphinx(); + let f = self.cxx; + unsafe { + // Find cell containing `x` + return ffi::eval_Function_f64(f, &expand_x); + } + } + } + //impl_diff!($n; 2 3 4 5 6 7 8); + )+ }; +} + +impl_mapping!(1 2 3); +// Extension to 1D should be easy. +//impl_diff!(1; 2 3 4 5 6 7 8); +impl_diff!(2; 2 3 4 5 6 7 8); + +impl<'py> GlobalAnalysis<f64, Bounds<f64>> for DolfinxPyFunction_f64<'py, 2, 2, 1> { + fn global_analysis(&self) -> Bounds<f64> { + let min = + unsafe { ffi::min_Function_f64_p2(self.cxx) }.map_or(f64::NEG_INFINITY, |res| res.v); + let max = unsafe { ffi::max_Function_f64_p2(self.cxx) }.map_or(f64::INFINITY, |res| res.v); + Bounds(min, max) + } +} + +/// Only second-order polynomials on a triangular mesh are implemented. +/// +/// May panic if assumptions of the underlying functionn do not hold despite +/// verification in the constructors of [`DolfinxPyFunction_f64`]. +impl<'py> MinMaxMapping<Loc<2>, f64> for DolfinxPyFunction_f64<'py, 2, 2, 1> { + fn minimise(&mut self, _tolerance: f64, _max_steps: usize) -> (Loc<2>, f64) { + let res = unsafe { ffi::min_Function_f64_p2(self.cxx) }.unwrap(); + (res.x.into(), res.v) + } + + fn maximise(&mut self, _tolerance: f64, _max_steps: usize) -> (Loc<2>, f64) { + let res = unsafe { ffi::max_Function_f64_p2(self.cxx) }.unwrap(); + (res.x.into(), res.v) + } +} + +/* +enum MaybeValid<A> { + Valid(A), + Invalid(A), +} + +use MaybeValid::*; + +impl<A> MaybeValid<A> { + #[inline] + fn uninitialised_ok(self) -> A { + match self { + Valid(a) => a, + Invalid(a) => a, + } + } +} +*/ + +impl<'py, const N: u32, const O: u32, const D: u32> DolfinxPyFunction_f64<'py, N, O, D> { + /// If owned, call `f_valid` with mutable data. + /// If not owned, create a new C++ fem.Function, uninitialised, and call `f_invalid` with its + /// mutable data, as well as the original immutable data. + #[inline] + fn with_mut_data_and_orig<F, G, Z>(&mut self, f_valid: F, f_invalid: G) -> Z + where + F: Fn(&mut [f64]) -> Z, + G: Fn(&mut [f64], &[f64]) -> Z, + { + if self.owned { + return f_valid(unsafe { ffi::data_mut_Function_f64(self.cxx) }); + } else { + // Create a completely new Function + let old_data = unsafe { ffi::data_mut_Function_f64(self.cxx) }; + let new = unsafe { ffi::similar_Function_f64(self.cxx) }; + let new_data = unsafe { ffi::data_mut_Function_f64(new) }; + let res = f_invalid(new_data, old_data); + // Invalidate Python reference; we're managing the object now + self.cxx = new; + self.u = None; + self.owned = true; + return res; + } + } + + /// If owned, call `f` with mutable data. + /// If not owned, create a new C++ fem.Function, uninitialised, and call `f` with its + /// mutable data. + #[inline] + fn with_mut_data<F, Z>(&mut self, f: F) -> Z + where + F: Fn(&mut [f64]) -> Z, + { + if self.owned { + return f(unsafe { ffi::data_mut_Function_f64(self.cxx) }); + } else { + // Create a completely new Function + let new = unsafe { ffi::similar_Function_f64(self.cxx) }; + let new_data = unsafe { ffi::data_mut_Function_f64(new) }; + let res = f(new_data); + // Invalidate Python reference; we're managing the object now + self.cxx = new; + self.u = None; + self.owned = true; + return res; + } + } + + /// Create a similar new function, and call `f` with both the mutable (unininitialised) + /// new data, and original data. + /// + /// Returns object for the new data. + #[inline] + fn with_new_data_and_orig<F>(&self, f: F) -> Self + where + F: Fn(&mut [f64], &[f64]), + { + // Create a completely new Function + let new = self.similar(); + let old_data = unsafe { ffi::data_Function_f64(self.cxx) }; + let new_data = unsafe { ffi::data_mut_Function_f64(new.cxx) }; + f(new_data, old_data); + new + } + + /// Similar to [`with_mut_data`], but check compatibility with `other` and also pass its + /// immutable data slice as the final argument. + #[inline] + fn with_mut_data_and_orig_binop<F, G, Z>( + &mut self, + other: &DolfinxPyFunction_f64<'py, N, O, D>, + f_valid: F, + f_invalid: G, + ) -> Z + where + F: Fn(&mut [f64], &[f64]) -> Z, + G: Fn(&mut [f64], &[f64], &[f64]) -> Z, + { + let data_other = unsafe { + ffi::check_compat_Function_f64(self.cxx, other.cxx).unwrap(); + ffi::data_Function_f64(other.cxx) + }; + self.with_mut_data_and_orig( + |data_self_valid| { + assert_eq!(data_self_valid.len(), data_other.len()); + f_valid(data_self_valid, data_other) + }, + |data_self_invalid, data_orig| { + assert_eq!(data_self_invalid.len(), data_other.len()); + assert_eq!(data_orig.len(), data_other.len()); + f_invalid(data_self_invalid, data_orig, data_other) + }, + ) + } + + /// Similar to [`with_new_data_and_orig`], but check compatibility with `other` and also + /// pass its immutable data slice as the final argument. + #[inline] + fn with_new_data_and_orig_binop<F>( + &self, + other: &DolfinxPyFunction_f64<'py, N, O, D>, + f: F, + ) -> Self + where + F: Fn(&mut [f64], &[f64], &[f64]), + { + let data_other = unsafe { + ffi::check_compat_Function_f64(self.cxx, other.cxx).unwrap(); + ffi::data_Function_f64(other.cxx) + }; + self.with_new_data_and_orig(|data, data_orig| { + assert_eq!(data.len(), data_other.len()); + assert_eq!(data_orig.len(), data_other.len()); + f(data, data_orig, data_other) + }) + } +} + +macro_rules! impl_unop { + ($trait:ident, $fn:ident) => { + impl<'py, const N: u32, const O: u32, const D: u32> $trait + for DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = Self; + + fn $fn(mut self) -> Self { + self.with_mut_data_and_orig( + |a| a.iter_mut().for_each(|x| *x = x.$fn()), + |a, a_orig| { + izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| { + *x = x_orig.$fn(); + }) + }, + ); + return self; + } + } + + impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait + for &'a DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = DolfinxPyFunction_f64<'py, N, O, D>; + + fn $fn(self) -> Self::Output { + self.with_new_data_and_orig(|a, a_orig| { + izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| { + *x = x_orig.$fn(); + }) + }) + } + } + }; +} + +use std::ops::Neg; + +impl_unop!(Neg, neg); + +macro_rules! impl_scalarop { + ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => { + impl<'py, const N: u32, const O: u32, const D: u32> $trait<f64> + for DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = Self; + + fn $fn(mut self, α: f64) -> Self { + self.with_mut_data_and_orig( + |a| { + a.iter_mut().for_each(|x| { + x.$fn_assign(α); + }) + }, + |a, a_orig| { + izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| { + *x = x_orig.$fn(α); + }) + }, + ); + return self; + } + } + + impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait<f64> + for &'a DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = DolfinxPyFunction_f64<'py, N, O, D>; + + fn $fn(self, α: f64) -> Self::Output { + self.with_new_data_and_orig(|a, a_orig| { + izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| { + *x = x_orig.$fn(α); + }) + }) + } + } + + impl<'py, const N: u32, const O: u32, const D: u32> $trait_assign<f64> + for DolfinxPyFunction_f64<'py, N, O, D> + { + fn $fn_assign(&mut self, α: f64) { + self.with_mut_data_and_orig( + |a| { + a.iter_mut().for_each(|x| { + x.$fn_assign(α); + }) + }, + |a, a_orig| { + izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| { + *x = x_orig.$fn(α); + }) + }, + ); + } + } + }; +} + +use std::ops::{Div, DivAssign, Mul, MulAssign}; + +impl_scalarop!(Mul, mul, MulAssign, mul_assign); +impl_scalarop!(Div, div, DivAssign, div_assign); + +macro_rules! impl_binop_consume { + ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => { + impl<'py, const N: u32, const O: u32, const D: u32> $trait<Self> + for DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = Self; + + #[inline] + fn $fn(self, other: Self) -> Self { + self.$fn(&other) + } + } + + impl<'py, const N: u32, const O: u32, const D: u32> $trait_assign<Self> + for DolfinxPyFunction_f64<'py, N, O, D> + { + #[inline] + fn $fn_assign(&mut self, other: Self) { + self.$fn_assign(&other) + } + } + + impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait<&'a Self> + for DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = Self; + + fn $fn(mut self, other: &'a Self) -> Self { + self.with_mut_data_and_orig_binop( + other, + |data, data_other| { + izip!(data.iter_mut(), data_other.iter()).for_each(|(x, y)| { + x.$fn_assign(y); + }) + }, + |data, data_orig, data_other| { + izip!(data.iter_mut(), data_orig.iter(), data_other.iter()).for_each( + |(x, x_orig, y)| { + *x = x_orig.$fn(y); + }, + ) + }, + ); + return self; + } + } + + impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait_assign<&'a Self> + for DolfinxPyFunction_f64<'py, N, O, D> + { + fn $fn_assign(&mut self, other: &'a Self) { + self.with_mut_data_and_orig_binop( + other, + |data, data_other| { + izip!(data.iter_mut(), data_other.iter()).for_each(|(x, y)| { + x.$fn_assign(y); + }) + }, + |data, data_orig, data_other| { + izip!(data.iter_mut(), data_orig.iter(), data_other.iter()).for_each( + |(x, x_orig, y)| { + *x = x_orig.$fn(y); + }, + ) + }, + ); + } + } + }; +} + +macro_rules! impl_binop_ref { + ($trait:ident, $fn:ident) => { + impl<'py, 'b, const N: u32, const O: u32, const D: u32> $trait<Self> + for &'b DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = DolfinxPyFunction_f64<'py, N, O, D>; + + fn $fn(self, other: Self) -> Self::Output { + self.$fn(&other) + } + } + + impl<'py, 'a, 'b, const N: u32, const O: u32, const D: u32> $trait<&'a Self> + for &'b DolfinxPyFunction_f64<'py, N, O, D> + { + type Output = DolfinxPyFunction_f64<'py, N, O, D>; + + fn $fn(self, other: &'a Self) -> Self::Output { + self.with_new_data_and_orig_binop(other, |data_dest, data, data_other| { + izip!(data_dest.iter_mut(), data.iter(), data_other.iter()).for_each( + |(x, y, z)| { + *x = y.$fn(z); + }, + ) + }) + } + } + }; +} + +macro_rules! impl_binop { + ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => { + impl_binop_consume!($trait, $fn, $trait_assign, $fn_assign); + impl_binop_ref!($trait, $fn); + }; +} + +use std::ops::{Add, AddAssign, Sub, SubAssign}; + +impl_binop!(Add, add, AddAssign, add_assign); +impl_binop!(Sub, sub, SubAssign, sub_assign); + +//use alg_tools::euclidean::Euclidean; +use alg_tools::linops::{VectorSpace, AXPY}; +//use alg_tools::norms::{Dist, HasDual, Norm, Normed, L2}; +use alg_tools::self_ownable; + +self_ownable!(DolfinxPyFunction_f64<'py, N, O, D> where 'py, const N: u32, const O: u32, const D: u32); + +impl<'py, const N: u32, const O: u32, const D: u32> Dist<L2, f64> + for DolfinxPyFunction_f64<'py, N, O, D> +{ + fn dist<I: Instance<Self>>(&self, other: I, p: L2) -> f64 { + other.eval_ref(|v| (self - v).norm(p)) + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> Norm<L2, f64> + for DolfinxPyFunction_f64<'py, N, O, D> +{ + fn norm(&self, _p: L2) -> f64 { + self.norm2_squared().sqrt() + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> VectorSpace + for DolfinxPyFunction_f64<'py, N, O, D> +{ + type Field = f64; + type PrincipalV = Self; + + /// Return a similar zero as `self`. + fn similar_origin(&self) -> Self { + let mut new = self.similar(); + new.set_zero(); + new + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> AXPY for DolfinxPyFunction_f64<'py, N, O, D> { + /// Computes `y = βy + αx`, where `y` is `Self`. + fn axpy<I: Instance<Self>>(&mut self, α: Self::Field, x: I, β: Self::Field) { + x.eval(|r| { + self.with_mut_data_and_orig_binop( + r, + |data, data_other| { + izip!(data.iter_mut(), data_other.iter()).for_each(|(y, x)| { + *y = β * (*y) + α * (*x); + }) + }, + |data, data_orig, data_other| { + izip!(data.iter_mut(), data_orig.iter(), data_other.iter()).for_each( + |(y, y_orig, x)| { + *y = β * (*y_orig) + α * (*x); + }, + ) + }, + ); + }) + } + + /// Set self to zero. + fn set_zero(&mut self) { + self.with_mut_data(|data| { + data.iter_mut().for_each(|y| *y = 0.0); + }) + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> HasDual + for DolfinxPyFunction_f64<'py, N, O, D> +{ + type DualSpace = Self; + + fn dual_origin(&self) -> <Self::DualSpace as VectorSpace>::PrincipalV { + self.similar_origin() + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> Normed for DolfinxPyFunction_f64<'py, N, O, D> { + type NormExp = L2; + + fn norm_exponent(&self) -> L2 { + // TODO: maybe H1 more approriate for O=2 + L2 + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> Euclidean + for DolfinxPyFunction_f64<'py, N, O, D> +{ + type PrincipalE = Self; + + fn dot<I: Instance<Self>>(&self, other: I) -> f64 { + let py = self.function_space.py(); + get_helpers(py) + .unwrap() + .dot + .call1(py, (self.own(), other.own())) + .unwrap() + .extract(py) + .unwrap() + /* + let data = unsafe { ffi::data_mut_Function_f64(self.cxx) }; + other.eval_decompose(|x| { + let data_other = unsafe { ffi::data_mut_Function_f64(x.cxx) }; + izip!(data.iter(), data_other.iter()) + .map(|(y, x)| y * x) + .sum() + })*/ + } + + fn norm2_squared(&self) -> f64 { + let py = self.function_space.py(); + get_helpers(py) + .unwrap() + .norm2_squared + .call1(py, (self.own(),)) + .unwrap() + .extract(py) + .unwrap() + } + + fn dist2_squared<I: Instance<Self>>(&self, other: I) -> f64 { + //other.eval_ref(|v| self.norm2_squared() + 2.0 * self.dot(v) + v.norm2_squared_div2()) + let py = self.function_space.py(); + get_helpers(py) + .unwrap() + .dist2_squared + .call1(py, (self.own(), other.own())) + .unwrap() + .extract(py) + .unwrap() + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/dolfinx_access/minmax_p2.cc Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,168 @@ +/// Minimisation code for second-order polynomial models from Fenics. +#include <array> +#include <cmath> + +#include <basix/finite-element.h> +#include <dolfinx/fem/Function.h> +#include <dolfinx/mesh/cell_types.h> +#include <nanobind/nanobind.h> +#include <nanobind/stl/array.h> // needed for array conversions in nanobind bindings. + +#include "dolfinx_access/minmax_p2.h" +#include "pointsource_pde/src/dolfinx_access.rs.h" + +using namespace dolfinx::fem; +namespace cell = basix::cell; +using namespace basix::element; +using namespace dolfinx_access; + +namespace dolfinx_access { + + // Since Fenics documentation is poor, and it does not seem to provide an easy way to access + // an internal P2 model (that would not involve vomotting your guts out 🤮 dealing with the + // putrid horror that C++ has become (and its documentation written by lawyers and bureacrats, + // which is even worse than that of Fenics), our idea here is to simply assume that f is a + // second-order polynomial on each cell, and then evaluate it on that cell at 6 points, to + // construct a new model based on our existing P2 model code that is also used for minimising + // arbitrary sums. + inline CoordValuePair minmax_Function_f64_p2(Function<double> const* f, bool max) { + auto v = f->x(); + // Mesh + auto fp = f->function_space(); + auto mesh = fp->mesh(); + auto geom = mesh->geometry(); + auto topo = mesh->topology(); + + // Check that we're working with P2 elements + // + // Fenics defaults to GLL warped midpoint placement; see + // https://docs.fenicsproject.org/dolfinx/v0.7.0.post0/python/demos/demo_lagrange_variants.html + // This is not a problem with our current implementation, which works with an + // second-order + // polynomials due to duplication of effort (see above, why), but if we would directly + // take the cell quadrature matrix from Fenics, we should restrict el.lagrange_variant() = + // lagrange_variant::equispaced. + auto el = fp->element()->basix_element(); + // printf("%d %d %d %d\n", el.degree(), el.cell_type(), el.lagrange_variant(), + // el.family()); + if (el.degree() != 2 || el.cell_type() != cell::type::triangle || + /*el.lagrange_variant() != lagrange_variant::equispaced ||*/ el.family() != + family::P) { + throw "Only equispaced Lagrange second-order polynomial elements are supported"; + } + if (cell_num_entities(mesh::CellType::triangle, 0) != 3) { + throw "A triangle should have three vertices"; + } + + // Check that we are dealing with a scalar function + if (fp->value_shape().size() != 0) { + throw "Only scalar functions are supported, obviously."; + } + + // Check that we're in two dimensions + if (topo->dim() != 2 || geom.dim() != 2) { + throw "Only two-dimensional meshes are supported"; + } + + // Check that there are no other types of eleemnts + auto entity_types = topo->entity_types(2); + for (auto& t : entity_types) { + if (t != mesh::CellType::triangle) { + throw "Only triangular meshes are supported"; + } + } + + auto adj = topo->connectivity(2, 0); + auto n_cells = adj->num_nodes(); + + // assert(n_cells == n_triangles); + + // printf("cell count? %d\n", n_cells); + // for (auto i = 0; i < n_cells; i++) { + // // assert(adj->num_links(i) == 3); // Should not need this + // auto d = fp->dofmap()->cell_dofs(i); + // printf("%zd\n", d.size()); + // } + + // DOF coordinates (may be more than nodes, e.g., intermediate points) + auto gx = geom.x(); + // auto num_points = gx.size() / 3; + // Map from cells to DOF indices + auto geom_dofmap = geom.dofmap(); + assert(geom_dofmap.extent(0) == (size_t)n_cells); + assert(geom_dofmap.extent(1) == 3); + auto fp_dofmap = fp->dofmap(); + + // auto dof_coords = fp->tabulate_dof_coordinates(); + + CoordValuePair res{{{0.0, 0.0}}, INFINITY}; + + // printf("%zd %zd\n", v->array().size(), gx.size()); + + for (auto i = 0; i < n_cells; i++) { + // Cell corners + auto j0 = geom_dofmap(i, 0); + auto j1 = geom_dofmap(i, 1); + auto j2 = geom_dofmap(i, 2); + + // Construct list of coordinates. Due to f->eval, we construct this is a single list + // with all the Fenics weirdness of hard-coded 3D. + std::array<double, 9> x = {gx[3 * j0], gx[3 * j0 + 1], 0, + gx[3 * j1], gx[3 * j1 + 1], 0, + gx[3 * j2], gx[3 * j2 + 1], 0}; + + // Find the solution on the simplex. + auto newres = minmax_dolfinx_p2_cell(f, x, i, max); + + // printf("%f [%f, %f] [%f %f %f %f %f %f] <%f, %f> <%f, %f> <‸%f, %f>\n", newres.v, + // newres.x[0], newres.x[1], values[0], values[1], values[2], values[3], + // values[4], values[5], x[0], x[1], x[3], x[4], x[6], x[7]); + if (newres.v < res.v) { + res = newres; + } + + // // Full DOF data + // auto d = fp_dofmap->cell_dofs(i); + // // TODO: array() gives just local part of vector + // assert(d.size() == 6); + // for (auto k = 0; k < d.size(); k++) { + // auto j = d[k]; + // // printf("%d %d\n", 3 * j + 1, nodes.size()); + // assert(2 * j + 1 <= nodes.size()); + // auto x = {nodes[3 * j], nodes[3 * j + 1]}; + // } + } + + // Negate result if maximising. + if (max) { + res.v = -res.v; + } + return res; + } + + CoordValuePair min_Function_f64_p2(Function<double> const* f) { + return minmax_Function_f64_p2(f, false); + } + + CoordValuePair max_Function_f64_p2(Function<double> const* f) { + return minmax_Function_f64_p2(f, true); + } + + // Create a Python module as well + NB_MODULE(_dolfinx_access, m) { + m.def("min_Function_f64_p2", &min_Function_f64_p2); + m.def("max_Function_f64_p2", &max_Function_f64_p2); + m.def("minmax_Function_f64_p2", &minmax_Function_f64_p2); + m.def("eval_Function_f64", &eval_Function_f64); + m.def("cell_Function_f64", &cell_Function_f64); + // m.def("cell_Mesh_f64", &cell_Mesh_f64); + m.def("cell_FunctionSpace_f64", &cell_FunctionSpace_f64); + nanobind::class_<CoordValuePair>(m, "CoordValuePair") + .def_rw("x", &CoordValuePair::x) + .def_rw("v", &CoordValuePair::v); + } + +} // namespace dolfinx_access + +// Forward declaration for manual registration +extern "C" PyObject* PyInit__dolfinx_access();
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/dolfinx_access/nanobind_helpers.cc Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,29 @@ +/// Helpers to transmit nanobind C++ → Python bindings Python → Rust → C++. +#include <Python.h> +#include <nanobind/nanobind.h> + +#include "pointsource_pde/include/dolfinx_access/nanobind_helpers.h" + +using namespace nanobind; + +namespace dolfinx_access { + bool check_Function_f64(const PyObject* obj) { + auto handle = nanobind::handle(obj); + return nanobind::isinstance<Function_f64>(handle); + } + + const Function_f64* cast_Function_f64(const PyObject* obj) { + auto handle = nanobind::handle(obj); + return nanobind::cast<const Function_f64*>(handle); + } + + Function_f64* cast_mut_Function_f64(PyObject* obj) { + auto handle = nanobind::handle(obj); + return nanobind::cast<Function_f64*>(handle); + } + + PyObject* wrap_Function_f64(Function_f64* f) { + return nanobind::cast(f).release().ptr(); + } + +} // namespace dolfinx_access
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/dolfinx_extras.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,54 @@ +# import weakref + +import numpy as np +import ufl +from dolfinx import fem +from dolfinx.fem.function import Function, FunctionSpace + +# _mass_matrices = weakref.WeakKeyDictionary() +_mass_matrices = {} + + +def get_mass_matrix(space: FunctionSpace): + # idex = space # Does not work due to fenicx being garbage + idx = id(space) # FIXME: This leaks memory 🤯 + if idx not in _mass_matrices: + u = ufl.TrialFunction(space) + v = ufl.TestFunction(space) + a = fem.form(ufl.inner(u, v) * ufl.dx) + m = fem.petsc.assemble_matrix(a) + m.assemble() + _mass_matrices[idx] = m + return _mass_matrices[idx] + + +def dot(f: Function, g: Function): + if g.function_space != f.function_space: + raise ValueError("Function spaces need to agree") + m = get_mass_matrix(f.function_space) + f_vec = f.x.petsc_vec + g_vec = g.x.petsc_vec + return g_vec.dot(m @ f_vec) + + +def norm2_squared(f: Function): + return dot(f, f) + + +def norm2(f: Function): + return norm2_squared(f).sqrt() + + +def dist2_squared(f: Function, g: Function): + if g.function_space != f.function_space: + raise ValueError("Function spaces need to agree") + m = get_mass_matrix(f.function_space) + f_vec = f.x.petsc_vec + g_vec = g.x.petsc_vec + mf = m @ f_vec + mg = m @ g_vec + return f_vec.dot(mf) + 2 * g_vec.dot(mg) + g_vec.dot(mf) + + +def dist2(f: Function, g: Function): + return np.sqrt(dist2_squared(f, g))
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/experiments.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,474 @@ +/*! +Access to experiments provided in Python code. +*/ + +use crate::dolfinx_access::DolfinxPyFunction_f64; +use crate::python_access::{ + process_error, Differentiable, HasProx, PythonMapping, PythonPlotFactory, +}; +use alg_tools::direct_product::Pair; +use alg_tools::error::{DynError, DynResult}; +use alg_tools::loc::Loc; +use anyhow::bail; +use clap::Parser; +use log::debug; +use pointsource_algs::measures::DiscreteMeasure; +use pointsource_algs::prox_penalty::RadonSquared; +use pointsource_algs::regularisation::{NonnegRadonRegTerm, RadonRegTerm, Regularisation}; +use pointsource_algs::run::{ + run_fb, run_fb_pair, AlgorithmConfig, DefaultAlgorithm, Named, RunError::NotImplemented, + RunnableExperiment, RunnableExperimentExtras, +}; +use pointsource_algs::{AlgorithmOverrides, CommandLineArgs, ExperimentSetup}; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyFunction}; +use serde::{Deserialize, Serialize}; +use serde_json; +use serde_with::skip_serializing_none; +use std::collections::HashMap; +use std::ffi::CString; +use std::path::Path; + +/// Command line experiment setup overrides +#[skip_serializing_none] +#[derive(Parser, Debug, Serialize, Deserialize, Default, Clone)] +#[pyclass(module = "pointsource_pde")] +pub struct Experiments { + /// List of Python setup files of experiments to perform + #[arg(value_name = "EXPERIMENT.PY")] + #[pyo3(get, set)] + experiments: Vec<String>, + + #[arg(long)] + #[pyo3(get, set)] + /// Regularisation parameter override. + /// + /// Only use if running just a single experiment, as different experiments have different + /// regularisation parameters. + alpha: Option<f64>, + + #[arg(long)] + #[pyo3(get, set)] + /// Noise variance override + variance: Option<f64>, +} + +/// An experiment implemented in Python code +#[derive(Serialize)] +pub struct PythonExperiment { + /// Name of the experiment + name: String, + /// Source file of experiment + filename: String, + /// The setup function + #[serde(skip_serializing)] + setup: Py<PyFunction>, + /// The Python module + #[serde(skip_serializing)] + #[allow(unused)] + module: Py<PyModule>, + /// Algorithm overrides + algorithm_overrides: HashMap<DefaultAlgorithm, AlgorithmOverrides<f64>>, +} + +impl ExperimentSetup for Experiments { + type FloatType = f64; + + fn runnables(&self) -> DynResult<Vec<Box<dyn RunnableExperiment<Self::FloatType>>>> { + debug!("Loading runnable experiments. Attaching Python."); + + Python::attach(|py| { + debug!("Loading json Python module."); + + // Load the "json" Python module, and "dumps" from it. + let dumps = PyModule::import(py, "json")?.getattr("dumps")?; + + // Load the Python file describing each experiment, and extract + // fields to file `PythonExperiment`. + self.experiments + .iter() + .cloned() + .map(|filename| -> DynResult<Box<dyn RunnableExperiment<_>>> { + debug!("Trying to load experiment from {filename}"); + let code = std::fs::read_to_string(&filename)?; + let modname = filename.as_bytes(); + // Add path of this file to module search path, to allow local dependencies + let parent = Path::new(&filename).parent().unwrap().to_str().unwrap(); + let locals = PyDict::new(py); + locals.set_item("this_path", parent)?; + py.run( + c_str!("import sys; sys.path.insert(0, this_path)"), + None, + Some(&locals), + )?; + // Load module + let module = PyModule::from_code( + py, + CString::new(code)?.as_ref(), + CString::new(filename.as_str())?.as_ref(), + CString::new(modname)?.as_ref(), + )?; + let name = module.getattr("name")?.extract()?; + let setup = module + .getattr("setup")? + .cast_into() // Check that a Pyfunction + .map_err(PyErr::from)? // Can't return a DowncastIntoError + .unbind(); + let algorithm_overrides = match module.getattr("algorithm_overrides") { + Err(_) => Default::default(), + Ok(o) => { + // This passes through JSON and Serde as pyo3 does not allow + // generics in AlgorithmOverides<_> - not even instantiating + // them to specific values. It would be interesting to write + // a proper serde Deserializer to not have to pass through JSON. + // One day. + serde_json::from_str(dumps.call1((o,))?.extract()?)? + } + }; + debug!("… Found {name} with algorithm overrides {algorithm_overrides:?}"); + Ok(Box::new(PythonExperiment { + name, + filename, + setup, + algorithm_overrides, + module: module.unbind(), + })) + }) + .try_collect() + }) + } +} + +/// Types for experiments in two dimensions with `f64` floats. +#[allow(non_camel_case_types)] +#[allow(unused)] +mod exp_f64_2 { + use super::*; + use crate::python_access::{DolfinxPyFunctionMarker, NumpyArrayMarker}; + + /// Type of unknowns + pub type Domain = DiscreteMeasure<Loc<2, f64>, f64>; + /// Type of derivatives of objective function + pub type DerivativeCodomain<'py> = DolfinxPyFunction_f64<'py, 2, 2, 1>; + //pub type AuxTODOSpecifyFlexiblise<'py> = DolfinxPyFunction_f64<'py, 2, 2, 1>; + pub type DerivativeMarker = DolfinxPyFunctionMarker<2, 2, 1>; + /// Type of data terms + pub type DataTerm<'py> = PythonMapping<'py, Domain, f64, Differentiable<DerivativeMarker>>; + /// Plotter + pub type PlotFactory<'py> = + PythonPlotFactory<'py, DerivativeCodomain<'py>, DerivativeCodomain<'py>>; +} + +/// Types for experiments in two dimensions with `f64` floats and aux +#[allow(non_camel_case_types, non_snake_case)] +#[allow(unused)] +mod exp_f64_2_auxvar { + use super::*; + use crate::python_access::{DolfinxPyFunctionMarker, NumpyArrayMarker}; + use numpy::Ix2; + + /// Type of unknowns + pub type Domain<'py> = Pair<DiscreteMeasure<Loc<2, f64>, f64>, AuxVar<'py>>; + /// Type of derivatives of objective function + pub type DerivativeCodomain<'py> = Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, AuxVar<'py>>; + /// Type of derivatives wrt. auxiliary variables. + pub type AuxVar<'py> = Pair< + DolfinxPyFunction_f64<'py, 2, 2, 1>, + Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, DolfinxPyFunction_f64<'py, 2, 2, 1>>, + >; + pub type DerivativeMarker = Pair< + DolfinxPyFunctionMarker<2, 2, 1>, + Pair< + DolfinxPyFunctionMarker<2, 2, 1>, + Pair<DolfinxPyFunctionMarker<2, 2, 1>, DolfinxPyFunctionMarker<2, 2, 1>>, + >, + >; + /// Type of data terms + pub type DataTerm<'py> = PythonMapping<'py, Domain<'py>, f64, Differentiable<DerivativeMarker>>; + /// Auxiliary objective + pub type AuxTerm<'py> = PythonMapping<'py, AuxVar<'py>, f64, HasProx>; + /// Plotter + pub type PlotFactory<'py> = + PythonPlotFactory<'py, DerivativeCodomain<'py>, DerivativeCodomain<'py>>; +} + +/// Types for experiments in two dimensions with `f64` floats and scalar aux +#[allow(non_camel_case_types, non_snake_case)] +#[allow(unused)] +mod exp_f64_2_auxvar_scalar { + use super::*; + use crate::python_access::{DolfinxPyFunctionMarker, NumpyArrayMarker}; + use numpy::Ix2; + + /// Type of unknowns + pub type Domain<'py> = Pair<DiscreteMeasure<Loc<2, f64>, f64>, AuxVar<'py>>; + /// Type of derivatives of objective function + pub type DerivativeCodomain<'py> = Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, AuxVar<'py>>; + /// Type of derivatives wrt. auxiliary variables. + pub type AuxVar<'py> = Pair<f64, Pair<f64, f64>>; + pub type DerivativeMarker = Pair<DolfinxPyFunctionMarker<2, 2, 1>, Pair<f64, Pair<f64, f64>>>; + /// Type of data terms + pub type DataTerm<'py> = PythonMapping<'py, Domain<'py>, f64, Differentiable<DerivativeMarker>>; + /// Auxiliary objective + pub type AuxTerm<'py> = PythonMapping<'py, AuxVar<'py>, f64, HasProx>; + /// Plotter + pub type PlotFactory<'py> = + PythonPlotFactory<'py, DerivativeCodomain<'py>, DerivativeCodomain<'py>>; +} + +#[pyclass(module = "pointsource_pde", name = "RegTerm")] +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)] +pub enum RegTermPy { + // Radon norm with weight `α`. + Radon(f64), + // Radon norm ith weight `α` and a positivity constraint. + NonnegRadon(f64), +} + +impl From<RegTermPy> for Regularisation { + fn from(reg: RegTermPy) -> Regularisation { + match reg { + RegTermPy::Radon(α) => Regularisation::Radon(α), + RegTermPy::NonnegRadon(α) => Regularisation::NonnegRadon(α), + } + } +} + +#[pyclass(module = "pointsource_pde")] +#[derive(Debug)] +pub struct Problem { + /// Data term. On the python side, this should be a `class` that implements + /// `apply` from [`DiscreteMeasure_2_f64`] (TODO: extended parameters) to floats, and `diff` + /// from the space space to [`DolfinxPyFunction_f64<2, 1, 2>`]. + #[pyo3(set)] + dataterm: Py<PyAny>, //exp_f64_2::DataTerm, + + // Regularisation + #[pyo3(set, get)] + regterm: RegTermPy, + + // Auxiliary variable regularisation term or similar + #[pyo3(set, get)] + auxterm: Option<Py<PyAny>>, + + // Initial auxiliary variable + #[pyo3(set, get)] + auxinit: Option<Py<PyAny>>, + + // Initial measure + #[pyo3(set, get)] + μinit: Option<DiscreteMeasure<Loc<2, f64>, f64>>, + + // Plotter + #[pyo3(set, get)] + plot_factory: Option<Py<PyAny>>, +} + +#[pymethods] +impl Problem { + #[new] + #[pyo3(signature = (dataterm, regterm, auxterm = None, auxinit = None, μinit = None, plot_factory=None))] + fn new( + dataterm: Py<PyAny>, + regterm: RegTermPy, + auxterm: Option<Py<PyAny>>, + auxinit: Option<Py<PyAny>>, + μinit: Option<DiscreteMeasure<Loc<2, f64>, f64>>, + plot_factory: Option<Py<PyAny>>, + ) -> Self { + Self { dataterm, regterm, auxterm, auxinit, μinit, plot_factory } + } + + #[getter] + fn get_dataterm<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> { + self.dataterm.bind(py).clone() + } +} + +impl RunnableExperiment<f64> for PythonExperiment { + fn name(&self) -> &str { + &self.name + } + + fn runall( + &self, + cli: &CommandLineArgs, + algs: Option<Vec<Named<AlgorithmConfig<f64>>>>, + ) -> DynError { + // Set up algorithms + let algorithms = algs.unwrap_or_else(|| vec![DefaultAlgorithm::FB.get_named()]); + + debug!( + "Entered PythonExperiment::runall for experimen {}. Attaching Python.", + self.name + ); + + Python::attach(|py| { + debug!("Calling Python-side experiment initialisation."); + + let prefix = self.start(cli)?; + + let problem: PyRef<Problem> = process_error( + "PythonExperiment::runall", + py, + self.setup + .call1(py, (&prefix,)) + .and_then(|r| r.extract(py).map_err(PyErr::from)), + )?; + + let save_extra = |_, ()| Ok(()); + let μ0 = problem.μinit.clone(); + + match (problem.regterm, &problem.auxterm, &problem.auxinit) { + (regterm, None, None) => { + debug!("… Extracting data term."); + let dataterm: exp_f64_2::DataTerm<'_> = problem.dataterm.extract(py)?; + debug!("… Extracting plotter."); + let plot_factory: exp_f64_2::PlotFactory<'_> = + if let Some(ref p) = problem.plot_factory { + p.extract(py)? + } else { + PythonPlotFactory::dummy() + }; + let make_plotter = |prefix| plot_factory.produce(prefix); + + debug!("Running."); + self.do_runall( + &prefix, + cli, + algorithms, + make_plotter, + save_extra, + μ0, + |p| match regterm { + RegTermPy::Radon(α) => { + run_fb(&dataterm, &RadonRegTerm(α), &RadonSquared, p, |_| { + Err(NotImplemented.into()) + }) + .map(|μ| (μ, ())) + } + RegTermPy::NonnegRadon(α) => { + run_fb(&dataterm, &NonnegRadonRegTerm(α), &RadonSquared, p, |_| { + Err(NotImplemented.into()) + }) + .map(|μ| (μ, ())) + } + }, + ) + } + (regterm, Some(auxterm), Some(z_)) => { + debug!("… Extracting auxiliary variable."); + let z0: PyResult<exp_f64_2_auxvar::AuxVar<'_>> = z_.extract(py); + match z0 { + Ok(z) => { + debug!("… Extracting data term."); + let dataterm: exp_f64_2_auxvar::DataTerm<'_> = + problem.dataterm.extract(py)?; + debug!("… Extracting plotter."); + let plot_factory: exp_f64_2::PlotFactory<'_> = + if let Some(ref p) = problem.plot_factory { + p.extract(py)? + } else { + PythonPlotFactory::dummy() + }; + let make_plotter = |prefix| plot_factory.produce(prefix); + debug!("… Extracting auxiliary term."); + let auxterm: exp_f64_2_auxvar::AuxTerm<'_> = auxterm.extract(py)?; + + debug!("Running."); + self.do_runall( + &prefix, + cli, + algorithms, + make_plotter, + save_extra, + (μ0, z), + |p| match regterm { + RegTermPy::Radon(α) => run_fb_pair( + &dataterm, + &RadonRegTerm(α), + &RadonSquared, + &auxterm, + p, + |_| Err(NotImplemented.into()), + ) + .map(|Pair(μ, _)| (μ, ())), + RegTermPy::NonnegRadon(α) => run_fb_pair( + &dataterm, + &NonnegRadonRegTerm(α), + &RadonSquared, + &auxterm, + p, + |_| Err(NotImplemented.into()), + ) + .map(|Pair(μ, _)| (μ, ())), + }, + ) + } + Err(_) => { + let z: exp_f64_2_auxvar_scalar::AuxVar<'_> = z_.extract(py)?; + debug!("… Extracting data term."); + let dataterm: exp_f64_2_auxvar_scalar::DataTerm<'_> = + problem.dataterm.extract(py)?; + debug!("… Extracting plotter."); + let plot_factory: exp_f64_2::PlotFactory<'_> = + if let Some(ref p) = problem.plot_factory { + p.extract(py)? + } else { + PythonPlotFactory::dummy() + }; + let make_plotter = |prefix| plot_factory.produce(prefix); + debug!("… Extracting auxiliary term."); + let auxterm: exp_f64_2_auxvar_scalar::AuxTerm<'_> = + auxterm.extract(py)?; + + debug!("Running."); + self.do_runall( + &prefix, + cli, + algorithms, + make_plotter, + save_extra, + (μ0, z), + |p| match regterm { + RegTermPy::Radon(α) => run_fb_pair( + &dataterm, + &RadonRegTerm(α), + &RadonSquared, + &auxterm, + p, + |_| Err(NotImplemented.into()), + ) + .map(|Pair(μ, _)| (μ, ())), + RegTermPy::NonnegRadon(α) => run_fb_pair( + &dataterm, + &NonnegRadonRegTerm(α), + &RadonSquared, + &auxterm, + p, + |_| Err(NotImplemented.into()), + ) + .map(|Pair(μ, _)| (μ, ())), + }, + ) + } + } + } + (_, _, _) => { + bail!("Errors in problem {} setup: auxiliary term or auxiliary variable initialisation given without the other", self.name); + } + } + })?; + + // Should run the experiment here, going through all algorithms. + Ok(()) + } + + fn algorithm_overrides(&self, alg: DefaultAlgorithm) -> AlgorithmOverrides<f64> { + self.algorithm_overrides + .get(&alg) + .cloned() + .unwrap_or(Default::default()) + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/full_sampling.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,97 @@ +import numpy as np +from dolfinx import fem +from petsc4py import PETSc +from pointsource_pde.dolfinx_extras import get_mass_matrix +from slepc4py import SLEPc + + +class FullSampling: + """ + A linear operator class for full sampling of a concentration (fem.Function) + """ + + def __init__( + self, + V, + ): + """ + M: number of laser beams + num_points: number of discrete points per beam + domain_size: [xmin, xmax, ymin, ymax] + """ + self.V = V + self.M = V.dofmap.index_map.size_global + A = get_mass_matrix(V) + solver = PETSc.KSP().create(A.comm) + solver.setOperators(A) + solver.setType(PETSc.KSP.Type.PREONLY) + solver.getPC().setType(PETSc.PC.Type.LU) + self.solver = solver + + # ‖x‖^2 = ‖Ax.array‖^2 + # ‖Rx‖^2 = ‖x.array‖^2 = ‖A^{-1}Ax.array‖^2 ≤ ‖A^{-1}‖‖x‖^2, + # so we need the inverse of the minimal eigenvalue. + E = SLEPc.EPS().create(A.comm) + E.setOperators(A) + E.setProblemType(SLEPc.EPS.ProblemType.NHEP) + E.setWhichEigenpairs(SLEPc.EPS.Which.SMALLEST_MAGNITUDE) + # E.setWhichEigenpairs(SLEPc.EPS.Which.LARGEST_MAGNITUDE) + E.setFromOptions() + E.solve() + + sigma = abs(E.getEigenvalue(0)) + self._opnorm = 1.0 / sigma + print("Full sampling opnorm = %f" % self._opnorm) + + def apply(self, x): + """ + This does not work with MPI + """ + return x.x.array + + def diff_adjdir(self, j, _x): + """ + This does not work with MPI + """ + # We need ⟨Rx,v⟩_{ℝ^n} = ⟨x,R^*v⟩_V + # We have ⟨x,R^*v⟩_V = ⟨Ax.array,[R^*v].array⟩_{ℝ^m} + # But Rx = x.array, so we need + # ⟨x.array,v⟩_{ℝ^n} = ⟨Ax.array,[R^*v].array⟩_{ℝ^m} + # Taking [R^*v].array=A^{-1}v, the conjugate works correctly. + tmp = fem.Function(self.V) + tmp.x.array[:] = j + tmp.x.scatter_forward() + return tmp + # TODO: is convection_diffusion actualy based on norms in ℝ^n? + # res = fem.Function(self.V) + # res.x.array[:] = 0.0 + # self.solver.solve(tmp.x.petsc_vec, res.x.petsc_vec) + # res.x.scatter_forward() + # return res + + def opnorm(self, *args): + # raise NotImplementedError("Need mesh weighting?") + return self._opnorm + + def lipschitz_factor(self, *args): + return self.opnorm(*args) + + def diff_adj_lipschitz_factor(self, *args): + return self.opnorm(*args) + + def diff_bound(self, *args): + return self.opnorm(*args) + + def bound(self, *args): + return self.opnorm(*args) + + # Construct observation noise + def noise(self, noise_level): + M = self.M + # Create diagonal covariance matrix Gamma_v_t of shape (M, M) + Gamma_v_t = np.diag((noise_level * np.random.rand(M)) ** 2) + # Use average variance from diagonal as scalar variance + std = np.sqrt(np.mean(np.diag(Gamma_v_t))) + # Generate noise with shape (M, 1) using scalar standard deviation + noise = np.random.normal(0, std, size=(M,)) + return noise
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/laser_sampling.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,181 @@ +import numpy as np +from dolfinx import fem, geometry +from mpi4py import MPI +from petsc4py import PETSc +from pointsource_pde.dolfinx_extras import get_mass_matrix +from scipy.sparse.linalg import svds +from slepc4py import SLEPc + +from dolfinx_access import cell_FunctionSpace_f64 + + +class LaserSampling: + """ + A linear operator class for laser sampling of a concentration (fem.Function) + """ + + def __init__( + self, + V, + domain, + nx, + ny, + beams, + num_segments=100, + domain_size=[-2, 2, -2, 2], + ): + self.V = V + self.domain = domain + self.nx, self.ny = nx, ny + self.M = len(beams) + self.domain_size = domain_size + self.num_segments = num_segments + + # Build R matrix (single function does everything) + self.R = self.matrix(beams) + + A = get_mass_matrix(V) + solver = PETSc.KSP().create(A.comm) + solver.setOperators(A) + solver.setType(PETSc.KSP.Type.PREONLY) + solver.getPC().setType(PETSc.PC.Type.LU) + self.solver = solver + + # ‖x‖^2 = ‖Ax.array‖^2 + # ‖Lx‖^2 = ‖Rx.array‖^2 = ‖RA^{-1}Ax.array‖^2 ≤ ‖RA^{-1}‖‖x‖^2 ≤ ‖R‖‖A^{-1}‖‖x‖^2, + # so we need the inverse of the minimal eigenvalue. + E = SLEPc.EPS().create(A.comm) + E.setOperators(A) + E.setProblemType(SLEPc.EPS.ProblemType.NHEP) + E.setWhichEigenpairs(SLEPc.EPS.Which.SMALLEST_MAGNITUDE) + # E.setWhichEigenpairs(SLEPc.EPS.Which.LARGEST_MAGNITUDE) + E.setFromOptions() + E.solve() + + sigma = abs(E.getEigenvalue(0)) + # self._opnorm = 1.0 / sigma + # print("Full sampling opnorm = %f" % self._opnorm) + + # SVD operator norm + _u, s, _v = svds(self.R, k=1) + self._opnorm = s[0] # / sigma + print(f"Laser sampling opnorm = {self._opnorm:.6f}") + + def matrix(self, beams): + """Combined: generate beams + segment + cell collision + dof weighting""" + dofmap = self.V.dofmap + tdim = self.domain.topology.dim + + # Geometry trees + + bb_tree = geometry.bb_tree(self.domain, tdim) + + n_dofs = self.V.dofmap.index_map.size_global + R = np.zeros((self.M, n_dofs)) + + xmin, xmax, ymin, ymax = self.domain_size + + for beam_idx, (start_pt, end_pt) in enumerate(beams): + # Segment into midpoints + vec = end_pt - start_pt + total_len = np.linalg.norm(vec) + seg_len = total_len / self.num_segments + + midpoints = [] + for i in range(self.num_segments): + t = (i + 0.5) / self.num_segments + midpoint = start_pt + t * vec + midpoints.append(midpoint) + # midpoints = np.array(midpoints, dtype=np.float64) # Force float64 + + # Ensure shape (N, 3), C-contiguous, read-only for DOLFINx + # assert midpoints.shape[1] == 2, ( + # "Expected 2D points, got shape[1]=2 but need padding to 3D" + # ) + # midpoints = np.pad( + # midpoints, ((0, 0), (0, 1)), mode="constant" + # ) # (N, 3) with z=0 + # midpoints = np.ascontiguousarray(midpoints) # C-order + # midpoints.setflags(write=False) # Read-only requirement + + # Find colliding cells + # cell_candidates = geometry.compute_collisions_points(bb_tree, midpoints) + # colliding_cells = geometry.compute_colliding_cells( + # self.domain, cell_candidates, midpoints + # ) + + # Cell→DOFs→weight accumulation + # cell_lists = list(colliding_cells.array) + # for seg_i, cell_idx_raw in enumerate(cell_lists): + # cell_idx = int(cell_idx_raw) # Convert np.int32 → Python int + for point in midpoints: + pad_point = np.array([point[0], point[1], 0.0]) + cell_idx = cell_FunctionSpace_f64(self.V._cpp_object, pad_point) + if cell_idx >= 0: # Valid cell index + # Get global DOFs + cell_dofs_local = dofmap.cell_dofs(cell_idx) + # Convert local dofs to numpy int32 array (DOLFINx requirement) + # local_dofs_array = np.asarray(cell_dofs_local, dtype=np.int32) + # local_dofs_array.setflags(write=False) # Read-only requirement + + # Batch convert all local indices to global in one call + cell_dofs_global = dofmap.index_map.local_to_global(cell_dofs_local) + cell_dofs = np.unique(cell_dofs_global) + + weight = seg_len / len(cell_dofs) + R[beam_idx, cell_dofs] += weight # ADD for overlaps! + + return R + + def apply(self, x): + """ + This does not work with MPI + """ + x.x.scatter_forward() + return self.R @ x.x.array + + def diff_adjdir(self, j, _x): + """ + This does not work with MPI + """ + # We need ⟨Rx,v⟩_{ℝ^n} = ⟨x,R^*v⟩_V + # We have ⟨x,R^*v⟩_V = ⟨Ax.array,[R^*v].array⟩_{ℝ^m} + # But Rx = R₀ x.array, so we need + # ⟨x.array,R₀^* v⟩_{ℝ^n} = ⟨Ax.array,[R^*v].array⟩_{ℝ^m} + # Taking [R^*v].array=A^{-1}R₀^* v, the conjugate works correctly. + tmp = fem.Function(self.V) + np.matmul(self.R.T, j, out=tmp.x.array) + tmp.x.scatter_forward() + return tmp + # TODO: is convection_diffusion actualy based on norms in ℝ^n? + # res = fem.Function(self.V) + # res.x.array[:] = 0.0 + # self.solver.solve(tmp.x.petsc_vec, res.x.petsc_vec) + # res.x.scatter_forward() + # return res + + def opnorm(self, *args): + # raise NotImplementedError("Need mesh weighting?") + return self._opnorm + + def lipschitz_factor(self, *args): + return self.opnorm(*args) + + def diff_adj_lipschitz_factor(self, *args): + return self.opnorm(*args) + + def diff_bound(self, *args): + return self.opnorm(*args) + + def codomain_bound(self, xbound=None): + if xbound is None: + raise Exception("Linear operators have unbounded range") + else: + return self.opnorm() * xbound + + # Construct observation noise + def noise(self, noise_level): + M = self.M + # Generate noise with shape (M, 1) using scalar standard deviation + noise = np.random.normal(0, noise_level, size=(M,)) + return noise
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/main.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,165 @@ +/*! +TODO: include README here. +*/ + +#![feature(maybe_uninit_array_assume_init)] +#![feature(iterator_try_collect)] +#![feature(once_cell_try)] +// We use unicode. We would like to use much more of it than Rust allows. +// Live with it. Embrace it. +#![allow(uncommon_codepoints)] +#![allow(mixed_script_confusables)] +#![allow(confusable_idents)] + +use alg_tools::error::DynResult; +use log::debug; +use measures::python::pymod as pymod_measures; +use pointsource_algs::common_main; +use pyo3::append_to_inittab; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use std::ffi::CStr; + +mod dolfinx_access; +mod experiments; +mod python_access; + +use experiments::Experiments; +use python_access::pymod_pointsource_pde; + +macro_rules! pymods { + [$($modname:expr),*] => {&[$( + (c_str!(concat!("pointsource_pde.", $modname)), + c_str!(concat!("src/", $modname, ".py")), + c_str!(include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/", $modname, ".py")))), + )*]}; +} + +const PY_MODULES: &[(&CStr, &CStr, &CStr)] = pymods![ + "dolfinx_extras", + "compose", + "measure", + "quadratic_dataterm", + "convection_diffusion", + "laser_sampling", + "full_sampling" +]; + +/// The entry point for the program. +pub fn main() -> DynResult<()> { + // Initialise logging + colog::init(); + + debug!("Appending Rust-side Python modules to inittab"); + + // Register the C++ module in Pythong. CURRENTLY NOT USED. + // Can't use DynResult above and ? here to get traces from Python. + dolfinx_access::register_python_ffi().unwrap(); + append_to_inittab!(pymod_measures); + append_to_inittab!(pymod_pointsource_pde); + + debug!("Initialising python"); + + // Initialise Python. + Python::initialize(); + + debug!("Loading Python modules from embedded source files"); + + // Load modules + Python::attach(|py| { + // Verify that there's no attempt to use MPI. + py.run( + c_str!( + r#" +import mpi4py + +mpi4py.rc.initialize = False # do not initialize MPI automatically +mpi4py.rc.thread_level = "multiple" +mpi4py.rc.threads = True + +from mpi4py import MPI + +MPI.Init() +size = MPI.COMM_SELF.Get_size() +if size != 1: + raise RuntimeError("MPI_COMM_SELF size must be 1, but is %d" % size) +size = MPI.COMM_WORLD.Get_size() +if size != 1: + raise RuntimeError("MPI_COMM_WORLD size must be 1, but is %d" % size) +"# + ), + None, + None, + )?; + + // FIXME: this is weird, besides registering in [`dolfinx_access::register_python_ffí`], + // we need to import our module into Python, once we have the GIL or otherwise + // nanobind::type<Function_f64>() will SIGSEGV. + let _dolfinx_access = PyModule::import(py, "dolfinx_access")?; + + //let _pointsource_pde = python_access::register_module(py); + + for &(modname, filename, filecontents) in PY_MODULES { + debug!("… Loading {}", modname.to_string_lossy()); + PyModule::from_code(py, filecontents, filename, modname)?; + } + + Ok::<_, PyErr>(()) + })?; + + debug!("Entering pointsource_algs::common_main"); + + common_main::<Experiments>() +} + +/* +fn py_main<'py>(py: Python<'py>, access_points: &PythonAccessPoints) -> PyResult<()> { + let main = PyModule::import(py, "__main__")?; + + //let fp = u_py.getattr("_V")?; + let mut u: DolfinxPyFunction_f64<2, 2, 1> = main.getattr("u")?.extract()?; + + let res = u.minimise(0.0, 0); + println!("min : {res:?}"); + + let resx = u.maximise(0.0, 0); + println!("max : {resx:?}"); + + let res2 = u.apply(Loc([0.5, 0.5])); + println!("apply : {res2:?}"); + + let res3 = u.differential(Loc([0.5, 0.5])); + println!("diff : {res3:?}"); + + // let dolfinx_fem = PyModule::import(py, "dolfinx.fem")?; + // let d_function = dolfinx_fem.getattr("Function")?; + // let d_function_new = d_function.getattr("__new__")?; + // let f = d_function_new.call1((d_function, fp))?; + + let mut u10 = &u * 10.0; + let obj = u10.into_pyobject(py)?; + + let locals = PyDict::new(py); + locals.set_item("obj", obj)?; + py.run( + c_str!( + "\ +import numpy as np +print(\"10 * min array \", np.min(obj.x.array)) +print(\"10 * max array \", np.max(obj.x.array))" + ), + None, + Some(&locals), + )?; + + Ok(()) +} +*/ + +// fn get_dolfinx_function_type_object(py: Python<'_>) -> PyResult<*mut pyo3::ffi::PyTypeObject> { +// // Import the module and get the class +// let fem = py.import("dolfinx.fem.function")?; +// let func_class: Bound<'_, pyo3::types::PyType> = fem.getattr("Function")?.downcast_into()?; +// // Convert to raw pointer +// Ok(func_class.as_type_ptr() as *mut pyo3::ffi::PyTypeObject) +// }
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/measure.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,10 @@ +import numpy as np + + +class Measure: + def __init__(self, points, alphas): + self.points = map(np.array, points) + self.alphas = alphas + + def __iter__(self): + return zip(self.points, self.alphas)
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/python_access.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,65 @@ +/*! +Wrappers to Python objects. +*/ + +use pyo3::types::{PyModule, PyModuleMethods, PyTracebackMethods}; +use pyo3::{pymodule, Bound, PyResult, Python}; + +use alg_tools::error::DynResult; +use anyhow::anyhow; + +mod diff_mapping; +mod function; +mod numpy_array; +mod plot; +mod prox_mapping; + +//mod measures; + +#[allow(unused_imports)] +pub use diff_mapping::{ + Basic, Differentiable, DolfinxPyFunctionMarker, NumpyArrayMarker, PythonMapping, +}; +//pub use pointsource_algs::measures::python::{DiscreteMeasure_1_f64, DiscreteMeasure_2_f64}; +// + +#[allow(unused_imports)] +pub use prox_mapping::HasProx; + +#[allow(unused_imports)] +pub use numpy_array::{NumpyArray_f64, NumpyMatrix_f64, NumpyVector_f64}; + +#[allow(unused_imports)] +pub use plot::PythonPlotFactory; + +/// Populates the Python module. +/// +/// Needs to be called with [`pyo3::append_to_inittab`]: +/// ``` +/// append_to_inittab!(pymod_pointsource_pde); +/// ``` +/// before initialising the intepreter. +#[pymodule] +#[pyo3(name = "pointsource_pde")] +pub(crate) fn pymod_pointsource_pde(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::<super::experiments::Experiments>()?; + m.add_class::<super::experiments::Problem>()?; + m.add_class::<super::experiments::RegTermPy>()?; + Ok(()) +} + +/// Converts a [`PyErr`] error in a [`Results`] error into a string with traceback. +pub fn process_error<'py, T>(fname: &str, py: Python<'py>, res: PyResult<T>) -> DynResult<T> +where +{ + match res { + Ok(res) => Ok(res), + Err(e) => Err(match e.traceback(py) { + None => anyhow!("Failed Python traceback in {fname}; original error {e}"), + Some(r) => match r.format() { + Err(e) => anyhow!("Failed Python traceback formatting in {fname}: {e}"), + Ok(o) => anyhow!("Python-side error in {fname}: {e}.\n{o}"), + }, + }), + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/python_access/diff_mapping.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,303 @@ +/*! +[`DifferentiableMapping`]s implemented in Python. +*/ +use super::{process_error, NumpyArray_f64}; +use crate::dolfinx_access::DolfinxPyFunction_f64; +use alg_tools::direct_product::Pair; +use alg_tools::error::DynResult; +use alg_tools::linops::IdOp; +use alg_tools::mapping::{ClosedSpace, DifferentiableImpl, Instance, Mapping, Space}; +use anyhow::anyhow; +use ndarray::Dimension; +use numpy::{Ix1, Ix2}; +use pointsource_algs::forward_model::{BoundedCurvature, BoundedCurvatureGuess}; +use pointsource_algs::prox_penalty::{RadonSquared, StepLengthBound, StepLengthBoundPair}; +use pyo3::conversion::FromPyObject; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::PyClass; +use std::marker::PhantomData; + +#[derive(Copy, Debug, Clone)] +/// Marker for differentiable PythonMappings +pub struct Differentiable<DerivativeDomainMarker>(DerivativeDomainMarker); + +#[derive(Copy, Debug, Clone)] +/// Marker for PythonMappings without further properties. +pub struct Basic; + +#[derive(Debug)] +pub struct PythonMapping<'py, Domain, Codomain, Marker> +where + Domain: Space, + Codomain: Space, +{ + pub(super) obj: Bound<'py, PyAny>, + pub(super) _phantoms: PhantomData<(Domain, Codomain, Marker)>, +} + +macro_rules! intern_many { + ($py:expr, $($method:literal),*) => {[ $( + intern!($py, $method) + ),*]} +} + +// Export for super. +pub(super) use intern_many; + +impl<'a, 'py, Domain, Codomain> FromPyObject<'a, 'py> + for PythonMapping<'py, Domain, Codomain, Basic> +where + Domain: Space, + Codomain: Space + PyClass + FromPyObject<'a, 'py>, +{ + type Error = PyErr; + + fn extract(obj_: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + let obj = obj_.to_owned(); + // Verify that the necessary methods exist + for method in intern_many!(obj.py(), "apply", "diff_lipschitz_factor") { + obj.getattr(method)?; //.downcast::<PyFunction>()?; + } + Ok(PythonMapping { obj, _phantoms: PhantomData }) + } +} + +impl<'a, 'py, Domain, Codomain, DerivativeMarker> FromPyObject<'a, 'py> + for PythonMapping<'py, Domain, Codomain, Differentiable<DerivativeMarker>> +where + Domain: Space, + Codomain: Space + FromPyObject<'a, 'py>, +{ + type Error = PyErr; + + fn extract(obj_: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + let obj = obj_.to_owned(); + // Verify that the necessary methods exist + for method in intern_many!(obj.py(), "apply", "diff", "diff_lipschitz_factor") { + obj.getattr(method)?; //.downcast::<PyFunction>()?; + } + Ok(PythonMapping { obj, _phantoms: PhantomData }) + } +} + +impl<'py, Domain, Codomain, AnyMarker> Mapping<Domain> + for PythonMapping<'py, Domain, Codomain, AnyMarker> +where + Domain: Space, + Domain::Principal: IntoPyObject<'py>, + Codomain: ClosedSpace + for<'a> FromPyObject<'a, 'py, Error = PyErr>, +{ + type Codomain = Codomain; + + /// Compute the value of `self` at `x`. + fn apply<I: Instance<Domain>>(&self, x: I) -> Self::Codomain { + // TODO: use references and internal mutability? + //x_py = x.own().to_python(py).unwrap(); + let apply = intern!(self.obj.py(), "apply"); + process_error( + "PythonMapping::apply", + self.obj.py(), + self.obj + .call_method1(apply, (x.own(),)) + .and_then(|r| r.extract()), + ) + .unwrap() + } +} + +impl<'py, Domain, Codomain, AnyMarker> PythonMapping<'py, Domain, Codomain, AnyMarker> +where + Domain: Space, + Codomain: Space + for<'a> FromPyObject<'a, 'py>, +{ + pub(crate) fn get_obj(&self) -> Bound<'py, PyAny> { + self.obj.clone() + } +} + +macro_rules! impl_differentiable { + ($derivative:ty, $marker:ty) => { + impl<'py, Domain> DifferentiableImpl<Domain> + for PythonMapping<'py, Domain, f64, Differentiable<$marker>> + where + Domain: Space, + Domain::Principal: IntoPyObject<'py>, + //$derivative: Space + for<'py> FromPyObject<'py>, + //f64 : for<'py> FromPyObject<'py>, + { + type Derivative = $derivative; + + /// Compute the value of `self` at `x`. + fn differential_impl<I: Instance<Domain>>(&self, x: I) -> $derivative { + // TODO: use references and internal mutability? + //x_py = x.own().to_python(py).unwrap(); + let diff = intern!(self.obj.py(), "diff"); + process_error( + "PythonMapping::differential_impl", + self.obj.py(), + self.obj + .call_method1(diff, (x.own(),)) + .and_then(|r| r.extract()), + ) + .unwrap() + } + + fn apply_and_differential_impl<I: Instance<Domain>>( + &self, + x: I, + ) -> (<Self as Mapping<Domain>>::Codomain, $derivative) { + // TODO: use references and internal mutability? + //x_py = x.own().to_python(py).unwrap(); + let apply_and_diff = intern!(self.obj.py(), "apply_and_diff"); + process_error( + "PythonMapping::apply_and_differential_impl", + self.obj.py(), + self.obj + .call_method1(apply_and_diff, (x.own(),)) + .and_then(|r| r.extract()), + ) + .unwrap() + } + } + }; +} + +impl<'py, Domain, Marker> BoundedCurvature<f64> + for PythonMapping<'py, Domain, f64, Differentiable<Marker>> +where + Domain: Space, +{ + fn curvature_bound_components( + &self, + _guess: BoundedCurvatureGuess, + ) -> (DynResult<f64>, DynResult<f64>) { + let m = intern!(self.obj.py(), "curvature_bound_components"); + match process_error::<(Option<f64>, Option<f64>)>( + "curvature_bound_components", + self.obj.py(), + self.obj.call_method0(m).and_then(|r| r.extract()), + ) { + Ok((l, θ2)) => ( + l.ok_or_else(|| anyhow!("l is None")), + θ2.ok_or_else(|| anyhow!("θ2 is None")), + ), + Err(e) => ( + Err(anyhow!("(same error as second component)")), + Err(e.into()), + ), + } + } +} + +impl<'py, Domain, Marker> + StepLengthBound<f64, PythonMapping<'py, Domain, f64, Differentiable<Marker>>> for RadonSquared +where + Domain: Space, +{ + fn step_length_bound( + &self, + f: &PythonMapping<'py, Domain, f64, Differentiable<Marker>>, + ) -> DynResult<f64> { + let m = intern!(f.obj.py(), "diff_lipschitz_factor"); + process_error( + "PythonMapping::diff_lipschitz_factor", + f.obj.py(), + f.obj.call_method0(m).and_then(|r| r.extract()), + ) + } +} + +impl<'py, 'a, Domain, Marker, Z> + StepLengthBoundPair<f64, PythonMapping<'py, Domain, f64, Differentiable<Marker>>> + for Pair<&'a RadonSquared, &'a IdOp<Z>> +where + Domain: Space, +{ + fn step_length_bound_pair( + &self, + f: &PythonMapping<'py, Domain, f64, Differentiable<Marker>>, + ) -> DynResult<(f64, f64)> { + let m = intern!(f.obj.py(), "diff_lipschitz_factor_pair"); + process_error( + "PythonMapping::diff_lipschitz_factor_pair", + f.obj.py(), + f.obj.call_method0(m).and_then(|r| r.extract()), + ) + } +} + +#[derive(Copy, Debug, Clone)] +/// This is a marker type for identifying the derivative codomain of a differentiable +/// [`PythonMapping`]. Ideally we would just use, e.g., +/// ``` +/// PythonMapping<Domain, f64, Differentiable<DolfinxPyFunction_f64<2,2,1>>> +/// ``` +/// to encode the derivative codomain into the type. However `DolfinxPyFunction_f64<2,2,1>` is +/// not [´Send`] so [`pyo3`] does not allow such a mapping in a `pyclass`. We don't actually +/// ever pass a `DolfinxPyFunction_f64` to Python—it's a Rust wrapper for a PyObject—but pyo3, +/// not knowing better, disallows the type from even appearing in signatures. (It's ok to +/// `pyo3` as `DifferentiableImpl::Domain`, though. That's why we need this marker in place +/// of the real type: +/// ``` +/// PythonMapping<Domain, f64, Differentiable<DolfinxPyFunctionMarker<2,2,1>>> +/// ``` +pub struct DolfinxPyFunctionMarker<const N: usize, const O: usize, const D: usize>; +pub struct NumpyArrayMarker<Ix: Dimension>(Ix); + +//impl_differentiable!(DolfinxPyFunction_f64<1,2,1>, DolfinxPyFunctionMarker<1,2,1>); +impl_differentiable!(DolfinxPyFunction_f64<'py, 2,2,1>, DolfinxPyFunctionMarker<2,2,1>); +impl_differentiable!( + Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, NumpyArray_f64<'py, Ix1>>, + Pair<DolfinxPyFunctionMarker<2, 2, 1>, NumpyArrayMarker<Ix1>> +); +impl_differentiable!( + Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, NumpyArray_f64<'py, Ix2>>, + Pair<DolfinxPyFunctionMarker<2, 2, 1>, NumpyArrayMarker<Ix2>> +); +impl_differentiable!( + Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, Pair<f64, Pair<f64, f64>>>, + Pair<DolfinxPyFunctionMarker<2, 2, 1>, Pair<f64, Pair<f64, f64>>> +); +impl_differentiable!( + Pair< + DolfinxPyFunction_f64<'py, 2, 2, 1>, + Pair<f64, Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, DolfinxPyFunction_f64<'py, 2, 2, 1>>>, + >, + Pair< + DolfinxPyFunctionMarker<2, 2, 1>, + Pair<f64, Pair<DolfinxPyFunctionMarker<2, 2, 1>, DolfinxPyFunctionMarker<2, 2, 1>>>, + > +); +impl_differentiable!( + Pair< + DolfinxPyFunction_f64<'py, 2, 2, 1>, + Pair< + NumpyArray_f64<'py, Ix2>, + Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, DolfinxPyFunction_f64<'py, 2, 2, 1>>, + >, + >, + Pair< + DolfinxPyFunctionMarker<2, 2, 1>, + Pair< + NumpyArrayMarker<Ix2>, + Pair<DolfinxPyFunctionMarker<2, 2, 1>, DolfinxPyFunctionMarker<2, 2, 1>>, + >, + > +); +impl_differentiable!( + Pair< + DolfinxPyFunction_f64<'py, 2, 2, 1>, + Pair< + DolfinxPyFunction_f64<'py, 2, 2, 1>, + Pair<DolfinxPyFunction_f64<'py, 2, 2, 1>, DolfinxPyFunction_f64<'py, 2, 2, 1>>, + >, + >, + Pair< + DolfinxPyFunctionMarker<2, 2, 1>, + Pair< + DolfinxPyFunctionMarker<2, 2, 1>, + Pair<DolfinxPyFunctionMarker<2, 2, 1>, DolfinxPyFunctionMarker<2, 2, 1>>, + >, + > +);
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/python_access/function.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,77 @@ +/*! +Conversions between [`DolfinxPyFunction_f64`] and `dolfinx.fem.Function` in Python. +*/ + +use crate::dolfinx_access::ffi; +use crate::dolfinx_access::DolfinxPyFunction_f64; +use pyo3::conversion::FromPyObject; +use pyo3::exceptions::PyException; +use pyo3::prelude::*; + +macro_rules! py_bail { + ($msg:literal $(,)?) => { + return Err(PyException::new_err(format!($msg))) + }; + ($err:expr $(,)?) => { + return Err(PyException::new_err(format!($err))) + }; + ($fmt:expr, $($arg:tt)*) => { + return Err(PyException::new_err(format!($fmt, $($arg)*))) + }; +} + +impl<'a, 'py, const N: u32, const O: u32, const D: u32> FromPyObject<'a, 'py> + for DolfinxPyFunction_f64<'py, N, O, D> +{ + type Error = PyErr; + + fn extract(u_: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + // We maintain our reference-counted copy + let u = u_.to_owned(); + // The "_cpp_object" attribute of a Python Dolfinx Function points to the C++ instance + let u_cpp = u.getattr("_cpp_object")?; + let cxx = unsafe { ffi::cast_mut_Function_f64(u_cpp.as_ptr() as *mut ffi::PyObject) } + .or_else(|err| py_bail!("CXX cast error: {}", err))?; + let info = unsafe { ffi::info_Function_f64(cxx) }; + if !info.triangular_mesh { + py_bail!("Triangular mesh required") + } + if info.order < O { + py_bail!("Insufficient order") + } + if info.codomain_dim != D { + py_bail!("Codomain of invalid size") + } + if info.domain_dim != N { + py_bail!("Domain of invalid size") + } + DolfinxPyFunction_f64::new_prechecked(u, cxx) + } +} + +// TODO: should probably use internal mutability to avoid just supporting mut references. +impl<'a, 'py, const N: u32, const O: u32, const D: u32> IntoPyObject<'py> + for &'a mut DolfinxPyFunction_f64<'py, N, O, D> +where + 'py: 'a, +{ + type Target = PyAny; + type Error = PyErr; + type Output = pyo3::Borrowed<'a, 'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> { + self._into_py_ref(py) + } +} + +impl<'py, const N: u32, const O: u32, const D: u32> IntoPyObject<'py> + for DolfinxPyFunction_f64<'py, N, O, D> +{ + type Target = PyAny; + type Error = PyErr; + type Output = pyo3::Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> { + self._into_py(py) + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/python_access/numpy_array.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,151 @@ +/*! +Python and C++ wrapper for Dolfinx Function<f64> +*/ + +use alg_tools::euclidean::wrap::{WrapGuard, WrapGuardMut, Wrapped}; +use alg_tools::types::Float; +use nalgebra::{DMatrix, DMatrixView, DMatrixViewMut, Dyn}; +use ndarray::{Dimension, Ix1, Ix2}; +use numpy::{PyArray, PyArrayMethods, PyReadonlyArray, PyReadwriteArray, ToPyArray}; +use pyo3::prelude::*; + +/// A helper structure of dealing with dolfinx functions. +/// `N` is the domain dimension, `O` the order, and `D` is the codomain dimension. +#[allow(non_camel_case_types)] +#[derive(Debug, Clone)] +pub enum NumpyArray_f64<'py, D> { + PyWrapped { + /// Python object. + x: Bound<'py, PyArray<f64, D>>, + }, + Nalgebra { + x: DMatrix<f64>, + }, +} + +#[allow(non_camel_case_types)] +#[allow(unused)] +pub type NumpyVector_f64<'py> = NumpyArray_f64<'py, Ix1>; + +#[allow(non_camel_case_types)] +#[allow(unused)] +pub type NumpyMatrix_f64<'py> = NumpyArray_f64<'py, Ix2>; + +impl<'a, 'py, D: Dimension> FromPyObject<'a, 'py> for NumpyArray_f64<'py, D> { + type Error = PyErr; + + fn extract(x: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + Ok(NumpyArray_f64::PyWrapped { x: x.to_owned().cast_into()? }) + } +} + +impl<'py> IntoPyObject<'py> for NumpyArray_f64<'py, Ix1> { + type Target = PyArray<f64, Ix1>; + type Error = PyErr; + type Output = pyo3::Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> { + match self { + NumpyArray_f64::PyWrapped { x } => x.into_pyobject(py).map_err(From::from), + NumpyArray_f64::Nalgebra { x } => x.to_pyarray(py).reshape((x.len(),)), + } + } +} + +impl<'py> IntoPyObject<'py> for NumpyArray_f64<'py, Ix2> { + type Target = PyArray<f64, Ix2>; + type Error = PyErr; + type Output = pyo3::Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> { + match self { + NumpyArray_f64::PyWrapped { x } => x.into_pyobject(py).map_err(From::from), + NumpyArray_f64::Nalgebra { x } => Ok(x.to_pyarray(py)), + } + } +} + +#[allow(non_camel_case_types)] +#[derive(Debug)] +pub enum ArrayGuard<'a, 'py, D: Dimension, F: numpy::Element + Float = f64> { + PyWrapped { + /// Python object. + x: PyReadonlyArray<'py, F, D>, + }, + Nalgebra { + x: &'a DMatrix<F>, + }, +} + +#[allow(non_camel_case_types)] +#[derive(Debug)] +pub enum ArrayGuardMut<'a, 'py, D: Dimension, F: numpy::Element + Float = f64> { + PyWrapped { + /// Python object. + x: PyReadwriteArray<'py, F, D>, + }, + Nalgebra { + x: &'a mut DMatrix<F>, + }, +} + +macro_rules! impl_euclidean { + ($dim:ty) => { + impl<'a, 'py> WrapGuard<'a, f64> for ArrayGuard<'a, 'py, $dim, f64> where 'py : 'a { + type View<'b> = DMatrixView<'b, f64, Dyn, Dyn> where Self : 'b; + + #[inline] + fn get_view(&self) -> Self::View<'_> { + match self { + ArrayGuard::PyWrapped{ref x} => x.as_matrix(), + ArrayGuard::Nalgebra{x} => x.as_view(), + } + } + } + + impl<'a, 'py> WrapGuardMut<'a, f64> for ArrayGuardMut<'a, 'py, $dim, f64> where 'py : 'a { + type ViewMut<'b> = DMatrixViewMut<'b, f64, Dyn, Dyn> where Self : 'b; + + #[inline] + fn get_view_mut(&mut self) -> Self::ViewMut<'_> { + match self { + ArrayGuardMut::PyWrapped{ref x} => x.as_matrix_mut(), + ArrayGuardMut::Nalgebra{x} => x.as_view_mut(), + } + } + } + + impl<'py> Wrapped for NumpyArray_f64<'py, $dim> where Self : 'py { + type WrappedField = f64; + type Guard<'a> = ArrayGuard<'a, 'py, $dim, f64> where Self : 'a; + type GuardMut<'a> = ArrayGuardMut<'a, 'py, $dim, f64> where Self : 'a; + type UnwrappedOutput = DMatrix<f64>; + type WrappedOutput = Self; + + #[inline] + fn get_guard(&self) -> Self::Guard<'_> { + match self { + NumpyArray_f64::PyWrapped{ ref x } => ArrayGuard::PyWrapped{ x : x.readonly() }, + NumpyArray_f64::Nalgebra{ ref x } => ArrayGuard::Nalgebra{ x }, + } + } + + #[inline] + fn get_guard_mut(&mut self) -> Self::GuardMut<'_> { + match self { + NumpyArray_f64::PyWrapped{ ref mut x } => ArrayGuardMut::PyWrapped{ x : x.readwrite() }, + NumpyArray_f64::Nalgebra{ ref mut x } => ArrayGuardMut::Nalgebra{ x }, + } + } + + + fn wrap(x: Self::UnwrappedOutput) -> Self { + NumpyArray_f64::Nalgebra{ x } + } + } + alg_tools::wrap!(f64; NumpyArray_f64<'py, $dim> where 'py); + }; +} + +impl_euclidean!(Ix1); +impl_euclidean!(Ix2);
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/python_access/plot.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,83 @@ +use super::process_error; +use alg_tools::types::Float; +use pointsource_algs::measures::RNDM; +use pointsource_algs::plot::Plotter; +use pyo3::conversion::FromPyObject; +use pyo3::intern; +use pyo3::prelude::*; +use std::marker::PhantomData; + +#[derive(Debug, Clone)] +pub struct PythonPlotFactory<'py, T1, T2> { + pub(super) obj: Option<Bound<'py, PyAny>>, + pub(super) _phantoms: PhantomData<(T1, T2)>, +} + +impl<'a, 'py, T1, T2> FromPyObject<'a, 'py> for PythonPlotFactory<'py, T1, T2> { + type Error = PyErr; + + fn extract(obj_: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + let obj = obj_.to_owned(); + // Verify that the necessary methods exist + obj.getattr(intern!(obj.py(), "produce"))?; + Ok(PythonPlotFactory { obj: Some(obj), _phantoms: PhantomData }) + } +} + +#[derive(Debug, Clone)] +pub struct PythonPlotter<'py, T1, T2> { + pub(super) obj: Option<Bound<'py, PyAny>>, + pub(super) _phantoms: PhantomData<(T1, T2)>, +} + +impl<'a, 'py, T1, T2> FromPyObject<'a, 'py> for PythonPlotter<'py, T1, T2> { + type Error = PyErr; + + fn extract(obj_: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + let obj = obj_.to_owned(); + // Verify that the necessary methods exist + obj.getattr(intern!(obj.py(), "plot"))?; + Ok(PythonPlotter { obj: Some(obj), _phantoms: PhantomData }) + } +} + +impl<'py, T1, T2, F, const N: usize> Plotter<T1, T2, RNDM<N, F>> for PythonPlotter<'py, T1, T2> +where + F: Float, + for<'a> &'a RNDM<N, F>: IntoPyObject<'py>, +{ + fn plot_spikes(&mut self, iter: usize, _g: Option<&T1>, _ω: Option<&T2>, μ: &RNDM<N, F>) { + if let Some(ref obj) = self.obj { + let plot = intern!(obj.py(), "plot"); + process_error( + "PythonPlotter::plot", + obj.py(), + obj.call_method1(plot, (iter, μ)).map(|_| ()), + ) + .unwrap() + } + } +} + +impl<'py, T1, T2> PythonPlotFactory<'py, T1, T2> { + pub fn dummy() -> Self { + PythonPlotFactory { obj: None, _phantoms: PhantomData } + } +} + +impl<'py, T1: Clone, T2: Clone> PythonPlotFactory<'py, T1, T2> { + pub fn produce(&self, prefix: String) -> PythonPlotter<'py, T1, T2> { + if let Some(ref obj) = self.obj { + let produce = intern!(obj.py(), "produce"); + process_error( + "PythonPlotFactory::produce", + obj.py(), + obj.call_method1(produce, (prefix,)) + .and_then(|r| r.extract()), + ) + .unwrap() + } else { + PythonPlotter { obj: None, _phantoms: PhantomData } + } + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/python_access/prox_mapping.rs Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,81 @@ +/*! +Implementation of [`Prox`] for mappings implemented in Python. +*/ + +use crate::python_access::process_error; + +use super::diff_mapping::intern_many; +use super::diff_mapping::PythonMapping; +use alg_tools::convex::Prox; +use alg_tools::mapping::{ClosedSpace, Instance, Mapping, Space}; +use pyo3::conversion::FromPyObject; +use pyo3::intern; +use pyo3::prelude::*; +use std::marker::PhantomData; + +#[derive(Copy, Debug, Clone)] +/// Marker for PythonMappings that implement Prox +pub struct HasProx; + +impl<'py, Domain> Prox<Domain> for PythonMapping<'py, Domain, f64, HasProx> +where + Domain: ClosedSpace + Clone + for<'a> FromPyObject<'a, 'py, Error = PyErr> + IntoPyObject<'py>, +{ + type Prox<'a> + = PythonProx<'py, Domain> + where + Self: 'a; + + fn prox_mapping(&self, τ: Self::Codomain) -> Self::Prox<'_> { + PythonProx { τ, obj: self.get_obj(), _phantoms: PhantomData } + } +} + +#[derive(Debug)] +pub struct PythonProx<'py, Domain> +where + Domain: Space + for<'a> FromPyObject<'a, 'py>, +{ + τ: f64, + obj: Bound<'py, PyAny>, + _phantoms: PhantomData<Domain>, +} + +impl<'py, Domain> Mapping<Domain> for PythonProx<'py, Domain> +where + Domain: ClosedSpace + for<'a> FromPyObject<'a, 'py, Error = PyErr> + IntoPyObject<'py>, +{ + type Codomain = Domain; + + /// Compute the value of `self` at `x`. + fn apply<I: Instance<Domain>>(&self, x: I) -> Self::Codomain { + // TODO: use references and internal mutability? + let apply = intern!(self.obj.py(), "prox"); + process_error( + "PythonProx::apply", + self.obj.py(), + self.obj + .call_method1(apply, (self.τ, x.own())) + .and_then(|r| r.extract()), + ) + .unwrap() + } +} + +impl<'a, 'py, Domain, Codomain> FromPyObject<'a, 'py> + for PythonMapping<'py, Domain, Codomain, HasProx> +where + Domain: Space, + Codomain: ClosedSpace + FromPyObject<'a, 'py>, +{ + type Error = PyErr; + + fn extract(obj_: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> { + let obj = obj_.to_owned(); + // Verify that the necessary methods exist + for method in intern_many!(obj.py(), "apply", "prox") { + obj.getattr(method)?; //.downcast::<PyFunction>()?; + } + Ok(PythonMapping { obj, _phantoms: PhantomData }) + } +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/quadratic_dataterm.py Thu Feb 26 09:32:12 2026 -0500 @@ -0,0 +1,53 @@ +import numpy as np +from numpy.linalg import norm + +""" +A quadratic data term +""" + + +class QuadraticDataTerm: + def __init__(self, opA, b, b_norm, xbound=None, λ=1.0): + self.opA = opA + self.b = b + self.b_norm = b_norm + self.xbound = xbound + self.λ = λ + + def apply(self, x): + v = self.opA.apply(x) - self.b + return norm(v) ** 2 * (self.λ / 2) + + def diff(self, x): + v = self.opA.apply(x) - self.b + j = self.λ * v + return self.opA.diff_adjdir(j, x) + + def apply_and_diff(self, x): + v = self.opA.apply(x) - self.b + val = norm(v) ** 2 * (self.λ / 2) + j = self.λ * v + d = self.opA.diff_adjdir(j, x) + return (val, d) + + def diff_lipschitz_factor(self, xbound=None): + if hasattr(self.opA, "opnorm"): + # Linear operator + return self.λ * self.opA.opnorm() ** 2 + else: + raise Exception("diff_lipschitz_factor for non nonlinear operators") + # ‖∇A(x)^*∇F(A(x)) - ∇A(y)^*∇F(A(y))‖ + # = ‖[∇A(x)^*-∇A(y)^*]∇F(A(x)) - ∇A(y)^*[A(y)-A(x)]‖ + # ≤ L_{∇A(x)}(M_A+M_b) + M_{∇A(y)^*} L_A. + la = self.opA.lipschitz_factor() + ma = self.opA.bound(xbound=xbound) + mb = self.b.norm() + mda = self.opA.diff_bound(xbound=xbound) + lda = self.opA.diff_adj_lipschitz_factor() + return self.λ * (lda * (ma + mb) + mda * la) + + # def diff_lipschitz_factor_pair(self, *args): + # return self.diff_lipschitz_factor() + + def diff_bound(self, xbound=None): + return self.λ * (self.opA.codomain_bound(xbound=xbound) + self.b_norm)
