diff --git a/AGENTS.md b/AGENTS.md index 8a5fb3e..5655079 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,7 +4,9 @@ Julia port of the Python `DiffusionGeometry` package — data-driven calculus, geometry, and topology on point clouds via heat diffusion and the carré du champ operator. -See `PORTING_PLAN.md` for scope and the phased plan. +The port is complete; development now diverges from upstream. See `README.md` for +scope and `docs/src/upstream-bugs.md` for where the Julia behaviour deliberately +differs from the Python. ## Agent skills diff --git a/PORTING_PLAN.md b/PORTING_PLAN.md deleted file mode 100644 index 5b03147..0000000 --- a/PORTING_PLAN.md +++ /dev/null @@ -1,447 +0,0 @@ -# Porting `DiffusionGeometry` (Python) → Julia: Scope & Phased Plan - -Source: `../DiffusionGeometry` (Python, ~8,350 LOC source + ~3,850 LOC tests). -Target: this package (`DiffusionGeometryJ`). - -Implements *Computing Diffusion Geometry* (Jones & Lanners, 2026): data-driven -calculus/geometry/topology on point clouds via heat diffusion and the carré du -champ operator. - -> **Progress (2026-07-10): the port is complete, including visualisation.** All six -> phases done, plus the graph/edge constructor path, the tensor differential-operator -> sugar and Hodge decompositions, ambient polyvectors (`to_ambient`), `wedge_operator`, -> the `block`/`hstack`/`vstack` block operators, and a Makie plotting extension -> (`dgplot` & friends + `dganimate`; `hodge_star_2_form` ported into the package proper -> and parity-tested). 404 tests green. Everything in the Python package is ported except -> the abstract `geometry_engine.py` (a base class whose body just sets `self.xp = numpy`). -> Three upstream bugs were found and corrected rather than replicated: two in -> `methods/geodesics.py` and the permutation-parity error in `basis_utils._perm_tables` -> (§8). See the status column in §5 and the progress log in §8. - ---- - -## 1. What the codebase actually is - -Two cleanly separable layers. - -### A. A numerical core of pure functions (~40% of LOC) -Takes arrays, returns arrays. No classes, no state. Ports almost mechanically. - -- `core/diffusion/`: `knn_graph`, `markov_chain`, `build_symmetric_kernel_matrix`, - `compute_eigenfunction_basis`, `carre_du_champ_knn`, `carre_du_champ_graph`, - `gamma_compound`, `gamma_02`, `gamma_02_sym`, `regularise_diffusion`, - `regularise_bandlimit`. -- `operators/differential_operators/`: `derivative_weak`, `hessian_*`, - `up_delta_weak`, `levi_civita_02_weak`, `lie_bracket_weak` — einsum-based - weak-form matrix builders. -- `utils/basis_utils.py`: combinatorial index/sign generation for wedge & - symmetric bases. - -### B. An OO orchestration layer (~60%) -Wires the core into a lazy, cached, operator-overloaded API. This is the real -design work — and where Julia's multiple dispatch produces a *cleaner* result -than the original. - -- `MarkovTriple` / `ImmersedMarkovTriple` — data container holding two - **callbacks** (`cdc`, `regularise`). -- `DiffusionGeometry` — the orchestrator, with `@cached_property` / `@lru_cache` - accessors (`grad`, `d(k)`, `laplacian(k)`, `hessian`, `levi_civita`, …). -- `DiffusionGeometryCache` — memoized γ-tensors. -- Tensor algebra: `Tensor` base + `Function` / `VectorField` / `Form` / - `Tensor02` / `Tensor02Sym` + their spaces + `DirectSum`. -- `LinearOperator` / `BilinearOperator` with `@` (compose), `.adjoint`, - `()` (apply), `.spectrum()`, `.inverse()`. - ---- - -## 2. Dependency mapping — the ones that need thought - -| Python | Julia | Note | -|---|---|---| -| `opt_einsum.contract` | **`OMEinsum.jl`** for 4–5 tensor contractions (path optimization like opt_einsum); **`Tullio.jl`** for simple batched ones | Biggest translation surface — 19 files use `contract`. **NB:** if targeting Reactant (§7), express hot kernels as reshape + batched `*` + broadcast, *not* Tullio (scalar loops don't trace). *Not yet a dependency:* Phase 2's only contractions (`carre_du_champ_knn`, `gamma_02`) were small enough to write as an explicit per-point loop / `mul!`. First real need is Phase 3's `derivative_weak` etc. | -| `sklearn.NearestNeighbors` | `NearestNeighbors.jl` (`KDTree` / `knn`, `sortres=true`) | Already **1-based**. Confirmed: distances match to ~1e-9, indices exact on a tie-free cloud. Tie-breaking may differ from sklearn, so Phase 2 parity feeds Python's kNN output into the downstream math rather than relying on it. | -| `scipy.sparse.linalg.eigsh` | `Arpack.eigs(; which=:LM)` or `KrylovKit.eigsolve` | **Caveat found:** `Arpack.eigs` does *not* return eigenvalues in ascending order like `eigsh`. `compute_eigenfunction_basis` sorts eigenvalues descending explicitly (φ₀ = Perron, λ≈1, first). Eigenvectors carry gauge freedom → parity tests compare eigen*values*. | -| `scipy.sparse.coo_matrix` | `SparseArrays.sparse(I,J,V)` | | -| `np.linalg.eigh` / `eig` | `LinearAlgebra.eigen` (`Hermitian` for `eigh`) | | -| `np.add.at` (scatter) | plain `for` loop with `+=` (order-independent) | `carre_du_champ_graph`, `regularise` | -| `scipy.special.comb`, `itertools.combinations` | `Combinatorics.jl` (`combinations`, `binomial`) | | -| `@cached_property` / `@lru_cache` | mutable struct nullable fields + memo `Dict` for parametric accessors (`d(k)`) | §3 | -| matplotlib / plotly (`visualisation.py`, 1353 LOC) | `Makie.jl` — **rewrite fresh, port last** | not a translation | - ---- - -## 3. The architecture remodel (the crux) - -The Python class tree collapses into Julia abstract types + structs + dispatch. - -### Markov triple — callbacks become function-typed fields -```julia -struct ImmersedMarkovTriple{C,R} - function_basis::Matrix{Float64} # (n, n0) - measure::Vector{Float64} # (n,) - cdc::C # (f,h) -> Array (a closure) - regularise::R # x -> x (closure or identity) - immersion_coords::Matrix{Float64} # (n, d) - data_matrix::Union{Nothing,Matrix{Float64}} - n::Int; n_function_basis::Int; dim::Int -end -``` -The `partial(...)` closures in `from_knn_kernel` become Julia closures directly — -a clean 1:1. - -### Lazy caching -Python `@cached_property` → mutable struct with nullable fields (or a small memo -helper). `@lru_cache` on `d(k)` / `laplacian(k)` → a `Dict{Int,LinearOperator}` -field. -```julia -mutable struct DiffusionGeometry - triple::ImmersedMarkovTriple - n_coefficients::Int; rcond::Float64 - cache::GammaCache - _grad::Union{Nothing,LinearOperator} - _d::Dict{Int,LinearOperator} - # ... -end -grad(dg) = @get! dg._grad build_grad(dg) # memoized accessor -d(dg, k) = get!(dg._d, k) do; build_d(dg, k) end -``` - -### Tensor algebra — dispatch replaces inheritance + numpy-fighting machinery -The `__array_ufunc__` / `__array_priority__` hooks exist *only* to make numpy -behave; Julia needs none of them. -```julia -abstract type AbstractTensor end -struct VectorField{S} <: AbstractTensor - space::S; coeffs::Array{ComplexF64}; batch_shape -end -Base.:+(a::T, b::T) where {T<:AbstractTensor} = wrap(a.space, a.coeffs .+ b.coeffs) -Base.:*(f::Function, t::AbstractTensor) = pointwise_product(f, t) # dispatch, not isinstance -⊼(a::Form, b::Form) = wedge(a, b) # `^` → a real operator -``` - -### Operators -- `L @ M` → `L ∘ M` (or `*`) -- `L.adjoint` → `adjoint(L)` / `L'` -- `L(x)` → make `LinearOperator` a functor: `(L::LinearOperator)(t) = ...` -- Space `__eq__` / `__hash__` (used as `lru_cache` / dict keys) → Julia `==` / `hash`. - ---- - -## 4. Cross-cutting hazards (ranked by bug-risk) - -1. **0-based → 1-based indexing.** Concentrated and dangerous in: `basis_utils` - (wedge / symmetric multi-indices, Laplace-expansion `children` / `signs`), - `gamma_compound` minors, `nbr_indices` gathers, `argmax` in `tune_kernel`, and - the `u /= u[0,0]` normalization. **Port `basis_utils` first and golden-test its - index arrays against Python before anything consumes them.** -2. **The big einsums.** e.g. `derivative_weak`: - `contract("pI,pJri,pJrj,p,r->IJij", ...)` (5 operands). opt_einsum optimizes - contraction order; naive Julia blows up memory. Use **OMEinsum** for these - (keeps the string ~verbatim and optimizes); reserve **Tullio** for simple - `pki,pkj,pk->pij` shapes. -3. **Memory layout.** numpy is row-major with the point axis `p` leading - (slowest). Julia is column-major. For a *faithful* first port, keep identical - index order and let OMEinsum/Tullio handle strides; only reorder axes (point - axis last) in a later optimization pass, gated by parity tests. -4. **Complex numbers.** `.conj()`, `real_if_close`, complex eigenreturns in - `LinearOperator.spectrum` / `inverse` — native in Julia; just don't force - `Float64` too early. -5. **Scatter (`np.add.at`)** is order-independent accumulation — a plain `for` - loop with `+=` is correct and fast. - ---- - -## 5. Phased plan - -Each phase ends at a **parity gate**: run the Python function on fixed random -input, save outputs (`.npy` / JLD2), assert Julia matches to `rtol=1e-8` -(numerics) / exact (index arrays). The existing **~3,850 lines of pytest** are the -spec — port the relevant tests alongside each phase. - -| Phase | Scope | Deliverable | Parity gate | Status | -|---|---|---|---|---| -| **0. Skeleton** | `Project.toml`, deps, CI, `pyparity/` harness that dumps Python reference outputs to disk | Package builds, `] test` runs empty suite | Harness produces reference fixtures | ✅ done | -| **1. Combinatorics + utils** | `basis_utils` (wedge/sym indices, signs, `kp1_children_and_signs`), `batch_utils`, `regularise` | Pure funcs | **Index arrays exactly match Python** (after +1 shift) | ✅ done (index fns, `regularise`, `batch_utils`, tensor-coeff `expand`/`symmetrise`, `form_to_ambient_polyvector`) | -| **2. Diffusion core** | `diffusion_process` (knn → markov → symmetric kernel → eigenbasis), `carre_du_champ_{knn,graph}`, `gamma_compound/02/02sym` | Build a `MarkovTriple` from a point cloud | γ-tensors match on a fixed torus sample | ✅ done (incl. `carre_du_champ_graph`) | -| **3. Weak-operator builders** | `derivative_weak`, `hessian_*`, `up_delta_weak`, `levi_civita_02_weak`, `lie_bracket_weak`, `metric_gram` | Pure matrix builders (the hard einsums) | Each weak matrix matches Python | ✅ done (OMEinsum `@ein_str`; row-major `np_reshape`) | -| **4. Spaces + tensor algebra** | `BaseTensorSpace` → concrete spaces, `Tensor` → concrete tensors, arithmetic/wedge/metric/`inner`/`norm`, basis conversions, `DirectSum` | `dg.function(x)`-style API | `g`, `inner`, pointwise products match | ✅ done (operator-coupled tensor methods deferred to Phase 5) | -| **5. Operators + orchestrator** | `LinearOperator` / `BilinearOperator` (`∘`, `'`, `spectrum`, `inverse`), `DiffusionGeometry` + `GammaCache`, all constructors (`from_point_cloud`, `from_edges`, …) | **Full public API**; README Quick Start runs | `grad`, `d(k)`, `laplacian(k).spectrum()`, `hessian`, `levi_civita`, curvature all match | ✅ done (incl. `from_edges`/`from_graph_kernel`/`from_sparse_matrix`) | -| **6. Methods + viz** | `methods/geodesics.py`, `methods/pde.py`; rewrite visualisation in Makie | End-to-end examples | Numeric methods parity-tested; plots reproduce notebook figures | ✅ done (Makie package extension: `dgplot`/recipes/`dganimate`; `hodge_star_2_form` parity-tested) | - -### Recommended tracer bullet -The vertical slice Phase 1 → 2 → minimal 4/5 needed to run -`laplacian(0).spectrum()` on a point cloud. That exercises -knn → markov → eigenbasis → cdc → weak matrix → gram → spectral solve — proving -every architectural seam end-to-end before fanning out to forms, Hessians, and -curvature. - ---- - -## 6. Effort & risk - -- **Phases 1–3** (pure numeric core): low risk, high mechanical throughput — the - bulk of correctness-critical code, with a clean oracle. -- **Phases 4–5** (the OO remodel): the design-heavy part. Risk is *architectural*, - not numerical — design the space/tensor/operator type hierarchy deliberately. -- **Phase 6**: the numeric methods (`pde.py`, `geodesics.py`) port cleanly; viz is - a rewrite, not a port — deferred/dropped. -- **Biggest single risks:** the multi-operand einsums in Phase 3 (memory blowup if - contracted naively) and off-by-one in Phase 1 combinatorics. Both are *contained* - and *golden-testable*, which is why they're front-loaded. - ---- - -## 7. Backend acceleration with Reactant.jl (optional, Phase 3.5) - -[Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) traces Julia functions -into MLIR/StableHLO and compiles them through **XLA**, giving one code path across -CPU / NVIDIA / AMD / Apple / TPU via `Reactant.set_default_backend(...)`, plus -Enzyme autodiff. It operates on `ConcreteRArray` / `TracedRArray` and — like XLA — -strongly prefers **dense arrays and static shapes**, tracing a *fixed* computation -graph. Data-dependent control flow, scalar indexing, dynamic shapes, and sparse / -iterative-solver library calls are where it fights you. - -**Verdict: use it surgically, not globally.** Good fit for the dense numerical -kernels; wrong tool for the rest. - -| Layer | Reactant fit | Why | -|---|---|---| -| Carré du champ contractions (Phase 2) | ✅ Strong | Dense batched contractions over the point axis — XLA's sweet spot; GPU/TPU-accelerable | -| Weak-form operator builders (Phase 3) | ✅ Strong | e.g. `contract("pI,pJri,pJrj,p,r->IJij", …)` — ideal dense XLA target | -| kNN graph, sparse kernel assembly, `eigsh`/Arpack | ❌ Poor | Sparse + KDTree + iterative eig aren't XLA-friendly and aren't traceable; run **once** at setup — no need to accelerate | -| Combinatorics / index arrays (Phase 1) | ❌ N/A | Dynamic, integer, tiny — keep as plain Julia | -| OO / caching / dispatch (Phases 4–5) | ❌ N/A | You don't trace mutable structs, dict memoization, or type dispatch — trace the *kernels*, keep orchestration in normal Julia | - -**Rules to keep it a drop-in choice, not an architectural commitment:** - -1. Build everything in **plain Julia first** (the phased plan as written). -2. Keep the numerical kernels behind a **small function boundary** and make them - **backend-parametric** (accept the array type / dispatch on it), so a - `ConcreteRArray` path can be added later without touching orchestration. -3. **Do not use Tullio for kernels you intend to accelerate** — its scalar-indexed - loops don't trace. Express those as reshape + batched `*` / `batched_mul` + - broadcast. -4. Only reach for Reactant if profiling shows the dense contractions dominate on - large point clouds (large `n`). For modest `n`, CPU BLAS is fine and the dense - spectral solves (LAPACK `eigen`) dominate — which Reactant doesn't help. -5. If you want GPU with less tracing friction, plain **`CUDA.jl`** on the - contraction kernels is often the lower-effort win; Reactant's edge is - *write-once → CPU/GPU/TPU + autodiff*. - -Treat this as an optional **Phase 3.5** acceleration pass, gated by the same -parity tests as Phase 3. - ---- - -## 8. Progress log - -- **Phase 0 — done.** Package skeleton (`Project.toml`, `src/DiffusionGeometryJ.jl`), - GitHub Actions CI, and the `pyparity/gen_fixtures.py` harness that dumps Python - reference outputs to `test/fixtures/*.npz`. Fixtures are committed so CI needs - no Python. -- **Phase 1 — mostly done.** `basis_utils` (`get_symmetric_basis_indices`, - `get_wedge_basis_indices`, `get_wedge_product_indices`, `kp1_children_and_signs`, - `lex_rank`) and `regularise` (`regularise_diffusion`, `regularise_bandlimit`) - ported and parity-tested. Index arrays match exactly under the +1 shift. - *Remaining:* `batch_utils`, `expand_symmetric_tensor_coeffs` / - `symmetrise_tensor_coeffs`; `form_to_ambient_polyvector` landed with the ambient block below. -- **Phase 2 — done.** `diffusion_process`, `carre_du_champ_knn`, - `gamma_compound/02/02sym`, and `(Immersed)MarkovTriple` + - `immersed_triple_from_point_cloud`. The γ-tensor gate passes on a torus sample - (`rtol 1e-8`); markov chain, symmetric kernel, and `K_sym` eigenvalues are also - parity-tested. *Deferred:* `carre_du_champ_graph` (edge/graph path — only - `from_edges` needs it). - -- **Phase 3 — done.** All weak-operator builders ported and parity-tested against - Python on a torus sample (`rtol 1e-8`): `derivative_weak` (k=0,1,2), - `hessian_functions`/`hessian_coords` (pointwise) + `hessian_02_weak` / - `hessian_02_sym_weak`, `up_delta_weak` (k=0,1, and the k≥2 Schur/`solve` branch), - `levi_civita_02_weak`, `lie_bracket_weak`, and `metric_gram`'s `gram` (+ `metric`). - The multi-operand einsums use **OMEinsum** (`@ein_str`, string kept ~verbatim from - the `opt_einsum` original). Key translation detail: the builders flatten their - multi-index tensors with numpy **C-order**, so a `np_reshape` helper - (`src/utils/reshape_utils.jl`) emulates row-major reshape (Julia's is - column-major) — every weak matrix is built with OMEinsum then `np_reshape`d. The - k≥2 up-Laplacian's batched `np.linalg.solve` is a per-`(p,J)` `lu!`/`ldiv!` loop. - Fixtures store the γ-tensor *inputs* + reference matrices so the gate targets the - weak builders in isolation (eigenbasis gauge / cdc are covered by earlier phases). - -- **Phase 4 — done.** The tensor algebra layer plus the minimal `DiffusionGeometry` - / `GammaCache` host needed to carry it. Ported: `batch_utils`, - `expand`/`symmetrise_tensor_coeffs`, `basis_conversions` (`_to`/`_from_pointwise_basis`), - `_metric_apply`, the abstract `AbstractTensorSpace`/`AbstractTensor`, all concrete - spaces + tensors (`FunctionSpace`/`ScalarFunction`, `VectorFieldSpace`/`VectorField`, - `FormSpace`/`Form`, `Tensor02Space`/`Tensor02`, `Tensor02SymSpace`/`Tensor02Sym`, - `DirectSumSpace`/`DirectSumElement`), arithmetic (`+ - *` scalar/function, - pointwise `/`, `^`), `wedge`, tensor product of 1-forms, `symmetrise` / `full_tensor` - / `transpose_tensor`, `flat` / `sharp`, `pack`/`unpack`, and the `g` / `inner` / - `l2_norm` / `pointwise_norm` metric API. The `tensor_algebra.npz` fixture builds a - Python `DiffusionGeometry` on the torus and stores its eigenbasis / measure / γ so - the Julia host is reconstructed gauge-for-gauge (the cache's `gamma_coords` is - seeded from the fixture because Python regularises the immersion coordinates - inside the constructor). Key translation notes: `Function` is renamed - `ScalarFunction` (Julia reserves `Function`); every flatten between a coefficient - vector and `(n1, C)`/`(n, C)` goes through `np_reshape` so the layout matches the - Gram basis; `_is_orthonormal` replicates numpy `allclose(atol=1e-10, rtol=1e-5)` - exactly so `_from_pointwise_basis` picks the same Gram-inverse branch. The - differential-operator accessors (`grad`, `d(k)`, `laplacian`, `hessian`, - `levi_civita`, `lie_bracket`) and the operator-coupled tensor methods - (`VectorField.operator`, `Tensor02` action, form interior product, `to_ambient`, - Hodge decomposition) are deferred to Phase 5 with the `LinearOperator` layer. - -- **Phase 5 — done.** The operator layer and the full orchestrator API. - `LinearOperator` (`src/operators/types/linear.jl`) carries weak/strong matrices - with lazy Gram conversion, `adjoint`/`'`, arithmetic, composition (`∘`/`*`), - the functor call `(L)(t)`, `spectrum`, and the spectral pseudo-`inverse`; - `BilinearOperator` (`bilinear.jl`) adds partial/full application and - `transpose`. The `DiffusionGeometry` accessors (`geometry_operators.jl`) wire the - Phase-3 weak builders + γ cache into operators — `grad`, `d(k)`, - `codifferential(k)`, `divergence`, `up_/down_/laplacian(k)`, `hessian`, - `lie_bracket`, `levi_civita` — memoised in `dg._op_cache`, plus - `riemann_curvature` / `sectional_curvature`. The deferred operator-coupled tensor - methods (`tensor_actions.jl`) landed too: `VectorField` directional-derivative - operator + `X(f)`, `Tensor02` operator/bilinear action + `to_ambient`, and the - `Tensor02Sym` delegates. The cache gained `hessian_functions` / `hessian_coords` - / `gamma_ambient`. `from_point_cloud` / `from_knn_graph` / `from_knn_kernel` - build a `DiffusionGeometry` end-to-end; the triple constructor now regularises - the immersion coordinates (mirroring Python's `resolve_immersion`) so the whole - pipeline matches gauge-for-gauge. Fixture `operators.npz` stores Python's - eigenbasis + regularised immersion coords so the Julia host reconstructs the same - `dg` and every operator *matrix* (not just gauge invariants) matches to - `rtol 1e-7`; `laplacian(k).spectrum()` eigenvalues and the curvatures match too. - Key translation notes: `∘`/`*` replace Python `@` for composition; `real_if_close` - replicates numpy's tolerance; the complex-spectrum sort mirrors - `lexsort((imag, real, |λ|))`; the weak builders take `cdc` positionally (not a - keyword). *Deferred at the time:* `from_edges`/`from_graph_kernel` (need - `carre_du_champ_graph`) and `VectorField.from_reconstruction` — both since landed. - -- **Phase 6 — done (numeric methods; viz dropped).** Ported the two numerical - methods to `src/methods/`. `solve_differential_operator` (`pde.jl`) diagonalises - an endomorphism and exponentiates it in its eigenbasis to evolve an initial - condition over `t_values`. Subtlety: the upstream reconstruction `Vᵀ diag(eᵗᵛ) V` - is **sensitive to each eigenvector's sign** (it inherits whatever sign LAPACK - returns, which isn't portable across BLAS); we canonicalise each eigenvector so - its largest-magnitude entry is positive (`_canonicalise_eigvec_signs!`) and the - fixture generator applies the same fix, so `ft_coeffs` matches to ~1e-14. - `geodesic_distances_function` (`geodesics.jl`) ports the cvxpy SOCP — maximise - the source coefficient of a correction subject to per-point 1-Lipschitz - constraints w.r.t. Γₓ (top-`dim` eigenpairs) — to **Convex.jl + SCS**. The - upstream reads `dg.cache.data_matrix` (nonexistent; it's `dg.triple.data_matrix`) - and `dg.cache.u` — both fixed in the port and in the inlined reference used to - generate the fixture (cvxpy was never even installed upstream, so the function - had never run). The SOCP optimum is effectively unique, so cvxpy/SCS and - Convex.jl/SCS agree to ~1e-12 (gated loosely at rtol/atol 1e-3 for cross-build - robustness). New deps: `Convex`, `SCS`, `Random`. Fixture `methods.npz`. - Visualisation (`visualisation.py`, pde's `gif_from_functions`) was deferred to a - Makie extension — see the 2026-07-10 entry. *Deferred at the time:* the graph path - (`carre_du_champ_graph` → `from_edges`) and `VectorField.from_reconstruction`. - -- **Graph / edge path (2026-07-08):** ported the last deferred constructor path. - `carre_du_champ_graph` (`carre_du_champ.jl`) — the edge version of the cdc: - per-edge weighted outer products scatter-added to target nodes (Python's - `np.add.at` → order-independent `+=` loops), both the mean-centred and - point-centred branches. `edge_index` is `(2, num_edges)`, **1-based** (row 1 - source `j`, row 2 target `i`). Wired into `immersed_triple_from_graph_kernel` / - `immersed_triple_from_edges` (`markov_triples.jl`) and the - `from_graph_kernel` / `from_edges` / `from_sparse_matrix` classmethods - (`diffusion_geometry.jl`). `from_edges` uses the row-stochastic kernel - `w_{ji} = 1/d(i)` with measure = in-degrees; `from_sparse_matrix` bridges a - `SparseMatrixCSC` (`findnz` → edges `(source=col, target=row)`). The graph path - carries **no regulariser / data_matrix** (matches Python) and defaults the - function basis to the `n×n` identity. Fixture `graph.npz` + `test/test_graph.jl` - (12 tests): the raw cdc (both branches), `from_graph_kernel` (measure, γ-coords, - Δ₀ weak + spectrum) and `from_edges` (in-degree measure, γ-coords, Δ₀ spectrum) - all match to rtol 1e-8. - -**Tensor sugar + Hodge: done (2026-07-09).** The differential-operator methods on -tensors (`grad`/`d`/`up_laplacian`/`laplacian`/`hessian` on `ScalarFunction`; -`d`/`codifferential`/`up_`/`down_`/`laplacian` on `Form`), the interior product -`(ω::Form)(X)` = `g(ω, flat(X))`, and `hodge_decomposition` for both -`ScalarFunction` (→ `(coexact_potential, harmonic)`) and `Form` (→ -`(exact_potential, coexact_potential, harmonic)`, with `coexact_potential === -nothing` at top degree). All live in `operators/tensor_actions.jl` and delegate to -the memoised `dg` accessors. Fixture `tensor_sugar.npz` reuses the `gen_operators` -torus so the Julia host rebuilds `dg` gauge-for-gauge; `test/test_tensor_sugar.jl` -(32 tests) checks the coefficients against Python at rtol 1e-7 *and* asserts the -decompositions reconstruct (`ω = dα + δβ + h`). - -- **Ambient polyvectors, wedge operator, block operators (2026-07-09):** the last - gaps. `form_to_ambient_polyvector` becomes `to_ambient` on `Form` (plus the - one-line `ScalarFunction` and `VectorField` cases), built from - `permutations_with_signs` (`utils/basis_utils.jl`) and an index-at-a-time raising - loop rather than a dynamically-constructed einsum string. `wedge_operator` and the - `block` / `hstack` / `vstack` constructors (`operators/types/direct_sum.jl`) join - them. Fixture `ambient.npz`, `test/test_ambient.jl` (52 tests, rtol 1e-7). Suite is - now **342 tests**. - - **Two upstream defects, corrected rather than replicated** (as with - `methods/geodesics.py` in Phase 6): - 1. `basis_utils._perm_tables` takes permutation parity over `np.tril_indices` - (pairs `i > j`), counting *concordant* pairs instead of inversions. Since - #concordant = k(k-1)/2 − #inversions, every sign picks up a constant - `(-1)^(k(k-1)/2)`: the identity permutation is assigned −1 for k = 2, 3, so - `form_to_ambient_polyvector` is globally sign-flipped for k ≡ 2, 3 (mod 4). - `gen_ambient` undoes the factor and asserts the bug still exists, so the - fixture fails loudly if upstream fixes it. `_perm_tables` has no other caller. - 2. `VectorField.from_reconstruction` is dead code — it reads - `dg.operators_engine.vector_field_to_quiver`, and neither the attribute nor - the map exists anywhere in the Python tree (it raises `AttributeError`). The - missing map is the Jacobian of `to_ambient`, - `A[(p,a),(k,i)] = Γ_ambient[p,a,i]·u[p,k]`; the Julia port builds it as - `vector_field_to_quiver` and least-squares inverts it. Having no reference to - compare against, it is gated by a round trip against `to_ambient` instead of a - fixture, and is reachable as `dg_vector_field(dg, data; mode=:reconstruct)`. - - `wedge_operator` returns a `LinearOperator` (the Python function returns the bare - coefficient matrix, which is `matrix(wedge_operator(a, l))`). - -- **Visualisation (2026-07-10):** ported `visualisation.py` as a Makie **package - extension** (`ext/DiffusionGeometryJMakieExt.jl`, loads with any Makie backend). The - public API is declared in `src/visualisation/api.jl` so the names export from the - package proper; `@recipe` in the extension adds methods to those stubs. `dgplot(t)` - dispatches on the tensor's type to the right recipe (`dgscatter`/`dgquiver`/`dg2form`/ - `dg3form`/`dgellipsoids`) — pulling points from `geometry(t)`, so it collapses the - Python `plot_x(dg.immersion_coords, t.to_ambient())` idiom to dispatch — and cleans the - axis it creates (`clean_fig`'s job). `dganimate` replaces `gif_from_functions` via - Makie's `record`. Camera projection / `overpic_labels` are Plotly/LaTeX workarounds - with no Makie analogue and are not ported. **`hodge_star_2_form`** — the one numerical - routine in `visualisation.py` (curl in notebook 5) — went into the package proper - (`src/visualisation/hodge_star.jl`) and is parity-tested (fixtures `hodge_star_d{2,3}_*`, - `test_visualisation.jl`). New weakdep `Makie`; test dep `CairoMakie`. Suite now 404 tests. - -**Remaining gaps:** none. Only `geometry_engine.py` (an abstract base whose body only -sets `self.xp = numpy`) is excluded. - -**Upstream bugs.** Four defects found in the Python reference over the course of the -port are catalogued — with reproductions, root causes and patches — in -[`docs/src/upstream-bugs.md`](docs/src/upstream-bugs.md). - -**Harness conventions established** (also in `README.md`): 1-based indexing with -the `v .+ 1` parity convention; `NPZ.jl` cannot read zero-element arrays, so the -generator skips empty fixtures and those edge cases are asserted directly in -Julia; gauge-free quantities (eigenvectors, kNN tie-breaking) are parity-tested -via invariants rather than the raw output. - ---- - -## Appendix: source module → target phase map - -| Python module | Phase | -|---|---| -| `utils/basis_utils.py`, `utils/batch_utils.py` | 1 | -| `core/diffusion/regularise.py` | 1 | -| `core/diffusion/diffusion_process.py` | 2 | -| `core/diffusion/carre_du_champ.py` | 2 | -| `core/diffusion/symmetric_kernel.py`, `markov_triples.py` | 2 | -| `operators/differential_operators/derivative.py` | 3 | -| `operators/differential_operators/hessian.py` | 3 | -| `operators/differential_operators/laplacian.py` | 3 | -| `operators/differential_operators/levi_civita.py` | 3 | -| `operators/differential_operators/lie_bracket.py` | 3 | -| `tensors/base_tensor/metric_gram.py` | 3 | -| `utils/basis_conversions.py` | 4 | -| `tensors/base_tensor/*`, `tensors/functions/*`, `tensors/vector_fields/*` | 4 | -| `tensors/forms/*`, `tensors/tensor02/*`, `tensors/tensor02sym/*` | 4 | -| `tensors/direct_sum/*` | 4 | -| `operators/types/{linear,bilinear,direct_sum}.py` | 5 | -| `core/geometry/{diffusion_geometry,cache,geometry_engine}.py` | 5 | -| `methods/geodesics.py`, `methods/pde.py` | 6 | -| `visualisation.py` | 6 (Makie extension; `hodge_star_2_form` in `src/`) | diff --git a/docs/src/upstream-bugs.md b/docs/src/upstream-bugs.md index 70065f2..cacee2d 100644 --- a/docs/src/upstream-bugs.md +++ b/docs/src/upstream-bugs.md @@ -11,7 +11,7 @@ Where a bug has a numerical parity target, the Julia port implements the *correc behaviour and `pyparity/gen_fixtures.py` stores a **corrected** reference, so the parity gate compares against the right answer rather than pinning the bug. Each correction is guarded: if upstream fixes the bug, fixture generation fails loudly -rather than silently double-correcting. See `PORTING_PLAN.md` §8. A test defect has +rather than silently double-correcting. A test defect has no parity target and so no such guard — §5 is a coverage hole, closed on the Julia side only. diff --git a/src/core/diffusion/carre_du_champ.jl b/src/core/diffusion/carre_du_champ.jl index efe7743..5e28e08 100644 --- a/src/core/diffusion/carre_du_champ.jl +++ b/src/core/diffusion/carre_du_champ.jl @@ -50,49 +50,167 @@ function carre_du_champ_knn(f::AbstractArray, h::AbstractArray, h_tail = size(h)[2:end] F = prod(f_tail; init=1) H = prod(h_tail; init=1) - f_flat = reshape(f, n, F) # 1-D tail ⇒ layout-safe - h_flat = reshape(h, n, H) + T = float(typeof(one(eltype(f)) * one(eltype(h)))) + + # The per-point work runs in a point-last layout: one point's values are + # contiguous, so the neighbour gather is a column copy rather than an + # F-strided walk over an (n, F) array, and the rank-k update accumulates into + # a dense H×F buffer instead of the strided `out[p, :, :]` view a point-first + # layout would force. Only that buffer is scattered back into the (n, F, H) + # result — its stride is n, but consecutive points land on the same cache + # lines, so the scatter streams. + ft = _point_last(f, n, F, T) # (F, n) + ht = f === h ? ft : _point_last(h, n, H, T) # (H, n) + + out = Array{T}(undef, n, F, H) # every entry is written below + + # Points are independent, so the loop splits into one task per chunk, each + # with its own scratch. Single-threaded this is the plain loop plus one task. + shared = f === h + chunks = _point_chunks(n, F * H * k) + if length(chunks) == 1 + _cdc_knn_chunk!(out, 1:n, ft, ht, shared, diffusion_kernel, nbr_indices, + bandwidths, use_mean_centres, F, H, k, T) + else + @sync for chunk in chunks + Threads.@spawn _cdc_knn_chunk!(out, chunk, ft, ht, shared, diffusion_kernel, + nbr_indices, bandwidths, use_mean_centres, + F, H, k, T) + end + end + + # 1/2ρ, applied in one pass. Kept out of the per-point loop because here the + # divisor varies along the contiguous axis, so the whole sweep vectorises. + scale = bandwidths === nothing ? 2.0 : reshape(2 .* bandwidths, n, 1, 1) + out ./= scale + return reshape(out, (n, f_tail..., h_tail...)) +end - cdc = zeros(eltype(float(one(eltype(f_flat)) * one(eltype(h_flat)))), n, F, H) - diff_f = Matrix{Float64}(undef, k, F) - diff_h = Matrix{Float64}(undef, k, H) - wdiff_h = similar(diff_h) +# Contiguous, roughly equal point ranges. Chunking by thread rather than by point +# keeps scratch allocation and task overhead O(nthreads). `work_per_point` is the +# rank-k update's flop count: small tensors (a scalar Γ, say) do too little per +# point to pay for a spawn, and splitting them costs more than it saves, so they +# fall back to a single in-line chunk. +const _MIN_WORK_PER_TASK = 1 << 20 + +function _point_chunks(n::Integer, work_per_point::Integer) + nt = min(Threads.nthreads(), n, max(1, (n * work_per_point) ÷ _MIN_WORK_PER_TASK)) + nt <= 1 && return (1:n,) + return [round(Int, (t - 1) * n / nt) + 1 : round(Int, t * n / nt) for t in 1:nt] +end - @inbounds for p in 1:n +function _cdc_knn_chunk!(out::Array{T,3}, points, ft, ht, shared::Bool, diffusion_kernel, + nbr_indices, bandwidths, use_mean_centres::Bool, + F::Int, H::Int, k::Int, ::Type{T}) where {T} + diff_f = Matrix{T}(undef, F, k) + diff_h = shared ? diff_f : Matrix{T}(undef, H, k) + wdiff_h = Matrix{T}(undef, H, k) + block = Matrix{T}(undef, H, F) + means_f = Vector{T}(undef, F) + means_h = shared ? means_f : Vector{T}(undef, H) + + @inbounds for p in points w = view(diffusion_kernel, p, :) nbrs = view(nbr_indices, p, :) - # neighbour values - for a in 1:F, j in 1:k - diff_f[j, a] = f_flat[nbrs[j], a] + _gather!(diff_f, ft, nbrs, F, k) + shared || _gather!(diff_h, ht, nbrs, H, k) + if use_mean_centres + _centre_mean!(diff_f, means_f, w, F, k) + shared || _centre_mean!(diff_h, means_h, w, H, k) + else + _centre_point!(diff_f, ft, p, F, k) + shared || _centre_point!(diff_h, ht, p, H, k) end - for b in 1:H, j in 1:k - diff_h[j, b] = h_flat[nbrs[j], b] + # blockᵀ = diff_h diag(w) diff_fᵀ; the 1/2ρ scaling comes after the loop + for j in 1:k + wj = w[j] + for b in 1:H; wdiff_h[b, j] = wj * diff_h[b, j]; end end - if use_mean_centres - for a in 1:F - m = 0.0 - for j in 1:k; m += w[j] * diff_f[j, a]; end - for j in 1:k; diff_f[j, a] -= m; end + _rank_k!(block, diff_f, wdiff_h, F, H, k) + for b in 1:H, a in 1:F + out[p, a, b] = block[b, a] + end + end + return out +end + +# (n, tail...) → a dense (prod(tail), n) working copy. +function _point_last(x::AbstractArray, n::Integer, M::Integer, ::Type{T}) where {T} + xt = Matrix{T}(undef, M, n) + xf = reshape(x, n, M) + @inbounds for a in 1:M, p in 1:n + xt[a, p] = xf[p, a] + end + return xt +end + +@inline function _gather!(dst, src, nbrs, M, k) + @inbounds for j in 1:k + q = nbrs[j] + for a in 1:M; dst[a, j] = src[a, q]; end + end +end + +# Subtract each row's kernel-weighted mean over the neighbourhood. +# +# `diff` is M×k, so the two axes want opposite loop orders and which one wins +# depends on the tensor. For a wide tensor the sweep runs with j outermost, which +# walks `diff` contiguously and vectorises over M. For a narrow one the whole M×k +# block is L1-resident anyway, so striding is free and it is cheaper to reduce +# each row into a scalar in two passes than to make three passes over M-vectors. +const _CENTRE_WIDE_M = 16 + +@inline function _centre_mean!(diff, means, w, M, k) + @inbounds if M >= _CENTRE_WIDE_M + @simd for a in 1:M; means[a] = zero(eltype(means)); end + for j in 1:k + wj = w[j] + @simd for a in 1:M + means[a] = muladd(wj, diff[a, j], means[a]) end - for b in 1:H - m = 0.0 - for j in 1:k; m += w[j] * diff_h[j, b]; end - for j in 1:k; diff_h[j, b] -= m; end + end + for j in 1:k + @simd for a in 1:M + diff[a, j] -= means[a] end - else - for a in 1:F, j in 1:k; diff_f[j, a] -= f_flat[p, a]; end - for b in 1:H, j in 1:k; diff_h[j, b] -= h_flat[p, b]; end end - # cdc[p] = diff_fᵀ diag(w) diff_h - for b in 1:H, j in 1:k; wdiff_h[j, b] = w[j] * diff_h[j, b]; end - @views mul!(cdc[p, :, :], transpose(diff_f), wdiff_h) + else + for a in 1:M + m = zero(eltype(diff)) + for j in 1:k; m = muladd(w[j], diff[a, j], m); end + for j in 1:k; diff[a, j] -= m; end + end end +end - scale = bandwidths === nothing ? fill(2.0, n) : 2 .* bandwidths - @inbounds for p in 1:n - cdc[p, :, :] ./= scale[p] +# blockᵀ = B Aᵀ (an H×F result) for the tall-thin A (F×k), B (H×k) of one +# neighbourhood. Below `_GEMM_MIN_WORK` a `mul!` would spend more time in BLAS +# call overhead than in arithmetic — and it is paid once per point — so small +# blocks accumulate rank-1 over j instead. Writing the transposed result puts the +# vectorised axis on H, which is the longer of the two for every tensor space here. +const _GEMM_MIN_WORK = 1 << 11 + +@inline function _rank_k!(blockT, A, B, F, H, k) + if F * H * k >= _GEMM_MIN_WORK + mul!(blockT, B, transpose(A)) + return + end + @inbounds begin + for a in 1:F, b in 1:H; blockT[b, a] = zero(eltype(blockT)); end + for j in 1:k, a in 1:F + t = A[a, j] + @simd for b in 1:H + blockT[b, a] = muladd(B[b, j], t, blockT[b, a]) + end + end + end +end + +# Subtract the value at the centre point itself. +@inline function _centre_point!(diff, src, p, M, k) + @inbounds for j in 1:k, a in 1:M + diff[a, j] -= src[a, p] end - return reshape(cdc, (n, f_tail..., h_tail...)) end """ diff --git a/src/core/diffusion/diffusion_process.jl b/src/core/diffusion/diffusion_process.jl index d7baea4..6615fc6 100644 --- a/src/core/diffusion/diffusion_process.jl +++ b/src/core/diffusion/diffusion_process.jl @@ -124,19 +124,32 @@ function tune_kernel(kernel_entries::AbstractMatrix, epsilons::AbstractVector) ne = length(epsilons) nk = length(kernel_entries) avg = Vector{Float64}(undef, ne) - @inbounds for e in 1:ne - inv_eps = 1.0 / epsilons[e] - s = 0.0 - for v in kernel_entries - s += exp(-v * inv_eps) - end - avg[e] = s / nk + entries = vec(kernel_entries) + # The sweep is `length(epsilons) * length(kernel_entries)` exp calls and is the + # bulk of `markov_chain`. Each ε is an independent reduction, so the sweep + # threads cleanly; the result does not depend on how it is split. + Threads.@threads for e in 1:ne + avg[e] = _kernel_sum(entries, epsilons[e]) / nk end criterion = diff(log.(avg)) ./ diff(log.(epsilons)) m = argmax(criterion) # first max, matches np.argmax return epsilons[m], 2 * criterion[m] end +# Σ exp(-v/ε). `exp(-x)` is exactly 0.0 once x passes 1075·log 2 ≈ 745.1, and +# adding an exact zero cannot change the running sum, so skipping those terms is +# bit-identical to summing all of them. It only bites at the small-ε end of the +# sweep — a few percent of the total — but it costs a compare against an exp. +function _kernel_sum(entries::AbstractVector{<:Real}, epsilon::Real) + inv_eps = 1.0 / epsilon + cutoff = 746.0 * epsilon + s = 0.0 + @inbounds for v in entries + v < cutoff && (s += exp(-v * inv_eps)) + end + return s +end + """ markov_chain(nbr_distances, nbr_indices; c=0, bandwidth_variability=-0.5, knn_bandwidth=8) -> (diffusion_kernel, bandwidths) @@ -182,10 +195,11 @@ function markov_chain(nbr_distances::AbstractMatrix, nbr_indices::AbstractMatrix @assert size(nbr_indices) == (n, knn_kernel) epsilons = 2 .^ collect(-10:0.25:9.75) + sq_distances = nbr_distances .^ 2 # used twice # 1. Kernel density estimate q0 with density kernel bandwidths ρ_A. bandwidths_A = compute_local_bandwidths(nbr_distances, knn_bandwidth, "l2") # (n,) - kernel_entries_A = nbr_distances .^ 2 ./ (bandwidths_A[nbr_indices] .* bandwidths_A) + kernel_entries_A = sq_distances ./ (bandwidths_A[nbr_indices] .* bandwidths_A) epsilon_A, dim_A = tune_kernel(kernel_entries_A, epsilons) kernel_A = exp.(-kernel_entries_A ./ epsilon_A) ./ ((π * epsilon_A) ^ (dim_A / 2)) density_estimate_A = vec(sum(kernel_A, dims=2)) ./ (n .* bandwidths_A .^ dim_A) # (n,) @@ -193,7 +207,7 @@ function markov_chain(nbr_distances::AbstractMatrix, nbr_indices::AbstractMatrix # 2. Kernel K with bandwidths ρ_B derived from q0. bandwidths_B = density_estimate_A .^ bandwidth_variability # (n,) bandwidths_B ./= median(bandwidths_B) - kernel_entries_B = nbr_distances .^ 2 ./ (bandwidths_B[nbr_indices] .* bandwidths_B) + kernel_entries_B = sq_distances ./ (bandwidths_B[nbr_indices] .* bandwidths_B) epsilon_B, dim_B = tune_kernel(kernel_entries_B, epsilons) kernel_B = exp.(-kernel_entries_B ./ epsilon_B) density_estimate_B = vec(sum(kernel_B, dims=2)) ./ (bandwidths_B .^ dim_B) # (n,)