src/laser_sampling.py

import basix
import numpy as np
from dolfinx import fem
from mpi4py import MPI
from petsc4py import PETSc
from pointsource_pde.dolfinx_extras import get_mass_matrix
from scipy.sparse import coo_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,
        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

        # Basix element matching V, used to tabulate phi_h exactly at points
        # self._element = self._create_basix_element()
        self._element = self.V.element.basix_element

        # 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}")

    # ------------------------------------------------------------------
    # Exact per-cell line integration
    #
    # A beam is split into the ordered sub-segments where it actually
    # crosses each element (via Cyrus-Beck clipping against the cell's
    # polygon), and each sub-segment is integrated exactly using the
    # 3-point Simpson rule (exact for phi_h of degree <= 2, since phi_h
    # pulled back along a straight line stays degree <= 2 in the
    # parameter t). This replaces the previous approach of sampling many
    # midpoints and splitting weight equally across a cell's dofs.
    # ------------------------------------------------------------------

    def _cell_vertices_2d(self, cell_idx):
        """Physical (x, y) vertices of a cell, ordered CCW around its centroid."""
        x_geom = self.domain.geometry.x
        verts = self.domain.geometry.dofmap[cell_idx]
        pts = x_geom[verts, :2]
        centroid = pts.mean(axis=0)
        angles = np.arctan2(pts[:, 1] - centroid[1], pts[:, 0] - centroid[0])
        return pts[np.argsort(angles)]

    def _clip_segment_to_cell(self, pts, x, y, eps=1e-12):
        """
        Cyrus-Beck clip of the segment x->y against the convex polygon `cell_idx`.
        Returns (t0, t1) in [0, 1] such that x + t*(y - x), t in [t0, t1], is the
        part of the segment inside the cell -- or None if there's no overlap.
        Works for any convex cell (triangle or quad); assumes straight-sided
        (affine, P1) geometry.
        """
        n_verts = len(pts)
        d = y - x
        t_lo, t_hi = 0.0, 1.0

        for i in range(n_verts):
            v0 = pts[i]
            v1 = pts[(i + 1) % n_verts]
            edge = v1 - v0
            normal = np.array([edge[1], -edge[0]])  # outward normal (CCW polygon)

            denom = normal @ d
            numer = normal @ (v0 - x)

            if abs(denom) < eps:
                if numer < -eps:  # segment parallel to edge, and entirely outside it
                    return None
                continue  # parallel and inside: this edge imposes no constraint

            t_edge = numer / denom
            if denom > 0:
                t_hi = min(t_hi, t_edge)
            else:
                t_lo = max(t_lo, t_edge)

            if t_lo > t_hi + eps:
                return None

        if t_lo > t_hi:
            return None
        return t_lo, t_hi

    def matrix(self, beams):
        """Combined: exact beam/cell intersection + exact per-cell line integration."""
        n_dofs = self.V.dofmap.index_map.size_global
        # R = np.zeros((self.M, n_dofs))
        rows, cols, data = [], [], []

        for cell_idx in range(len(self.domain.geometry.dofmap)):
            pts = self._cell_vertices_2d(cell_idx)
            for beam_idx, (start_pt, end_pt) in enumerate(beams):
                clip = self._clip_segment_to_cell(pts, start_pt, end_pt)
                if clip is not None:
                    t0, t1 = clip
                    if t1 - t0 > 1e-10:
                        beam_v = end_pt - start_pt
                        x0 = start_pt + t0 * beam_v
                        x1 = start_pt + t1 * beam_v
                        midpt = start_pt + (0.5 * (t1 + t0)) * beam_v
                        """
                        Exact int_{x0}^{x1} phi_h ds for phi_h in V restricted to `cell_idx`, as a
                        (global_dofs, row_values) pair such that row_values @ phi_h.x.array[global_dofs]
                        equals the integral. Uses the 3-point Simpson rule (endpoints + midpoint),
                        which is exact since phi_h has degree <= 2 along the line.
                        """
                        x_geom = self.domain.geometry.x
                        verts = self.domain.geometry.dofmap[cell_idx]
                        v0, v1, v2 = (
                            x_geom[verts[0], :2],
                            x_geom[verts[1], :2],
                            x_geom[verts[2], :2],
                        )
                        J = np.column_stack([v1 - v0, v2 - v0])
                        """Map a physical point p to reference coordinates on an affine triangle."""
                        pts_ref = np.array(
                            [
                                np.linalg.solve(J, x0 - v0),
                                np.linalg.solve(J, midpt - v0),
                                np.linalg.solve(J, x1 - v0),
                            ]
                        )
                        tab = self._element.tabulate(0, pts_ref)[
                            0
                        ]  # (npoints, ndof_local, value_size)
                        tab = tab[
                            ..., 0
                        ]  # scalar element -> drop trailing value_size=1 axis

                        length = (t1 - t0) * np.linalg.norm(beam_v)
                        weights = np.array([1.0, 4.0, 1.0]) * length / 6.0
                        local_dofs = self.V.dofmap.cell_dofs(cell_idx)
                        # global_dofs = self.V.dofmap.index_map.local_to_global(local_dofs)
                        # R[beam_idx, local_dofs] += weights @ tab
                        rows.append(np.full(len(local_dofs), beam_idx, dtype=np.int64))
                        cols.append(np.asarray(local_dofs, dtype=np.int64))
                        data.append(weights @ tab)

        # return R

        if rows:
            rows = np.concatenate(rows)
            cols = np.concatenate(cols)
            data = np.concatenate(data)
        else:
            rows = np.array([], dtype=np.int64)
            cols = np.array([], dtype=np.int64)
            data = np.array([], dtype=np.float64)

        return coo_matrix((data, (rows, cols)), shape=(self.M, n_dofs)).tocsr()

    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)
        res = fem.Function(self.V)
        tmp.x.array[:] = self.R.T @ j
        tmp.x.scatter_forward()
        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_chain_lipschitz_factor(self, *args):
        return self.opnorm(*args)

    def diff_bound(self, *args):
        return self.opnorm(*args)

    # Construct observation noise
    def noise(self, noise_level, rng=None):
        if rng is None:
            rng = np.random.default_rng()
        M = self.M
        # Generate noise with shape (M, 1) using scalar standard deviation
        noise = rng.normal(0, noise_level, size=(M,))
        return noise