experiments/laser_and_mirrors_aux.py

#
# 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,
    ScaleAuxStepHack,
    SumOfSeparableFunctions,
)
from pointsource_pde.convection_diffusion import (
    BoxedQuadraticRegularisation,
    ConvectionDiffusion,
    PlotFactory,
    XBound,
)
from pointsource_pde.laser_sampling import LaserSampling
from pointsource_pde.quadratic_dataterm import QuadraticDataTerm

from dolfinx_access import eval_Function_f64

# Give name to the problem
name = "laser_and_mirrors_aux"

# Fine-tune algorithms
alg_finetune = {
    "tau0": 0.02,
    "sigmap0": 0.0001,
    "tautheta": 0.07,
    "gradv_lipest_mult": 1,
    "transport_tolerance": 1e5,
    "inner_method": "Exact",
    "merge": True,
    "fitness_merging": True,
    "merge_radius": 0.1,
    "merge_interp": False,
    "scaling_heuristic": None,
    "extra_weight_optimisation_steps": None,
}

algorithm_overrides = {
    "RadonFB": alg_finetune,
    "RadonSlidingFB": alg_finetune,
    "RadonPDPS": alg_finetune,
    "RadonSlidingPDPS": alg_finetune,
}

reg_params = {
    "alpha": 2e-3,
    "alpha_k": 60,
    "alpha_c": 0.06,
    "c_noise": 0.02,
    "k_noise": 0.02,
    "seed": 2393,
    "true_init": False,
    "true_target": False,
    "noise": 0.01,
    "sim_oversample": 2,
    "sim_oversample_time": 2,
    "scale_cstep": 300,
    "do_non_sliding": True,
}


# Setup shared between experiment with known/unknown conductivity and diffusivity
def generic_setup(
    prefix,
    simple=True,
    rng=None,
    μ_true=DiscreteMeasure_2_f64(
        [([0.1, 0.3], 0.08), ([0.4, 0.25], 0.05), ([0.25, 0.13], 0.06)]
    ),
    k=0.01,
    r=0.5,
    θ=30 * np.pi / 180,
    c_noise=reg_params["c_noise"],
    k_noise=reg_params["k_noise"],
    noise=reg_params["noise"],
    sim_oversample=reg_params["sim_oversample"],
    sim_oversample_time=reg_params["sim_oversample_time"],
    scale_cstep=reg_params["scale_cstep"],
):

    u0 = lambda x: 0.0 * x[1]
    g = lambda x: 0.0 * x[1]
    domain_dims = [[0, 0.5], [0, 0.5]]
    num_steps = 50
    res = 32

    # Create convection-diffusion PDE
    pde = ConvectionDiffusion(
        u0=u0,
        g=g,
        domain_dims=domain_dims,
        num_steps=num_steps,
        nx=res,
        ny=res,
    )

    # Create convection-diffusion PDE for simulation, to avoid inverse crimes
    pde_sim = ConvectionDiffusion(
        u0=u0,
        g=g,
        domain_dims=domain_dims,
        # num_steps=(sim_oversample-1) * (num_steps + 1) + num_steps,
        num_steps=sim_oversample_time * (num_steps - 1) + 1,
        nx=res * sim_oversample,
        ny=res * sim_oversample,
    )

    # 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_dims[0][0],
        pde.domain_dims[0][1],
        pde.domain_dims[1][0],
        pde.domain_dims[1][1],
    )
    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)
    L_i_sim = LaserSampling(pde_sim.V, pde_sim.domain, pde_sim.nx, pde_sim.ny, beams)

    v0 = 1.0

    if simple:
        # Scalar convectivity and diffusivity
        c1 = r * np.cos(θ)
        c2 = r * np.sin(θ)
        # k = 0.001  # high wind
        print("True aux variables: k = %f, c1 = %f, c2 = %f" % (k, c1, c2))
        k_sim = k
        c1_sim = c1
        c2_sim = c2
    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

        c1_sim = fem.Function(pde_sim.V)
        c1_sim.interpolate(velocity_x)
        c2_sim = fem.Function(pde_sim.V)
        c2_sim.interpolate(velocity_y)
        k_sim = fem.Function(pde_sim.V)
        k0 = 1.0
        k_sim.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_sim = pde_sim.apply((μ_true, (k_sim, (c1_sim, c2_sim))))
    u_n_list = pde.apply((μ_true, (k, (c1, c2))))

    b_true = [L_i_sim.apply(u) for u in u_n_list_sim[0::sim_oversample_time]]

    # Add noise
    level_k = abs(k)
    level_c = r
    k_noisy = max(0.001, relnoise(k, k_noise, rng, level=level_k))
    c1_noisy, c2_noisy = relnoise((c1, c2), c_noise, rng, level=level_c)
    aux_noisy = (k_noisy, (c1_noisy, c2_noisy))
    b = [
        b_i + L_i_sim.noise(noise * norm(b_i) / np.sqrt(np.size(b_i)), rng)
        for b_i in b_true
    ]

    # Estimate bounds on domains and values
    dat_bound = 2 * max([norm(b_i) for b_i in b])
    μ_bound = 1.3 * sum([alpha for _point, alpha in μ_true.iter_padded()])
    pde.xbound = XBound(mu_dual=μ_bound, k_min=k * 0.9, k_max=k * 1.1, c_Linf=r * 0.9)

    # 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,
    )
    if scale_cstep is not None:
        dat = ComposeFnWithOperator(dat, ScaleAuxStepHack((1.0, scale_cstep)))

    dat.curvature_bound_components = lambda: (None, None, None)

    # Save true and initial data
    plot_factory = PlotFactory(pde)
    plotter = plot_factory.produce(prefix)
    ω, _ = dat.diff((DiscreteMeasure_2_f64([]), auxtrue))
    ω_true, (dk, (dc1, dc2)) = 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,
        k_true=k,
        c1_true=c1,
        c2_true=c2,
        k_noisy=k_noisy,
        c1_noisy=c1_noisy,
        c2_noisy=c2_noisy,
        true_u_n_list_array=[u_n_list[i].x.array for i in plotter.plot_frames],
    )

    delta = 0.001 * k
    dk_est = (
        dat.apply((μ_true, (k + delta, (c1, c2)))) - dat.apply((μ_true, auxtrue))
    ) / delta
    print("k-diff-test:", dk, ", ", dk_est)
    delta = 0.001 * c1
    dc1_est = (
        dat.apply((μ_true, (k, (c1 + delta, c2)))) - dat.apply((μ_true, auxtrue))
    ) / delta
    print("c1-diff-test:", dc1, ", ", dc1_est)
    delta = 0.001
    μ_perturb = DiscreteMeasure_2_f64(
        [([0.1, 0.3], 0.08 + delta), ([0.4, 0.25], 0.05), ([0.25, 0.13], 0.06)]
    )
    dcμ_est = (dat.apply((μ_perturb, auxtrue)) - dat.apply((μ_true, auxtrue))) / delta
    print(
        "μ-diff-test:",
        eval_Function_f64(ω_true._cpp_object, [0.1, 0.3, 0.0]),
        ", ",
        dcμ_est,
    )

    return dat, auxtrue, aux_noisy, μ_bound, μ_true, plot_factory, pde


def relnoise(v, σ, rng, level=None):
    if isinstance(v, float):
        return float(rng.normal(v, σ * (abs(v) if level is None else level)))
    else:
        return map(
            float, rng.normal(v, σ * (np.linalg.norm(v) if level is None else level))
        )


# Setup the problem
def setup(
    prefix,
    alpha=reg_params["alpha"],
    alpha_k=reg_params["alpha_k"],
    alpha_c=reg_params["alpha_c"],
    c_noise=reg_params["c_noise"],
    k_noise=reg_params["k_noise"],
    noise=reg_params["noise"],
    seed=reg_params["seed"],
    sim_oversample=reg_params["sim_oversample"],
    sim_oversample_time=reg_params["sim_oversample_time"],
    true_init=False,
    true_target=False,
    scale_cstep=reg_params["scale_cstep"],
):
    rng = np.random.default_rng(seed=seed)

    dat, auxtrue, aux0, _μ_bound, μ_true, plot_factory, pde = generic_setup(
        prefix,
        rng=rng,
        c_noise=c_noise,
        k_noise=k_noise,
        noise=noise,
        sim_oversample=sim_oversample,
        sim_oversample_time=sim_oversample_time,
        scale_cstep=scale_cstep,
    )

    # Override Lipschitz parameter to use direct step length choice
    dat.override_lipschitz = 1.0
    dat.override_lipschitz_pair = (0.0, 0.0)

    reg = RegTerm.NonnegRadon(alpha)

    μ0 = DiscreteMeasure_2_f64([])

    print("aux init ", aux0)
    aux = BoxedQuadraticRegularisation(
        (0.005, -1.0),
        (1.0, 1.0),
        (alpha_k, alpha_c),
        auxtrue if true_target else aux0,
        (1, scale_cstep) if scale_cstep is not None else None,
    )
    print("Initial value:", dat.apply((μ0, aux0)) + aux.apply(aux0))
    print("Value at true μ:", dat.apply((μ_true, auxtrue)) + aux.apply(auxtrue))

    return Problem(
        dat,
        reg,
        aux,
        auxtrue if true_init else aux0,
        μ0,
        plot_factory=plot_factory,
        domain=pde.domain_dims,
    )