diff --git a/examples/paper/pet_simulations_comparison_block_b.py b/examples/paper/pet_simulations_comparison_block_b.py new file mode 100644 index 00000000..5456c94c --- /dev/null +++ b/examples/paper/pet_simulations_comparison_block_b.py @@ -0,0 +1,894 @@ +"""Comparisons of block b closed for dicom and vtu images. + +""" + +from pathlib import Path + +import matplotlib.cm as cm +import matplotlib.pyplot as plt +import numpy as np + +import darsia + +# ! ---- Constants + +cm2m = 1e-2 # conversion cm -> m +h2s = 60**2 # conversion hours -> seconds +ml_per_hour_to_m3_per_s = cm2m**3 / h2s + +# ! ---- Model paramaters + +porosity_2d = 0.2321 +fracture_aperture = 0.1 * cm2m +depth = 1.95 * cm2m +injection_rate = 15 * ml_per_hour_to_m3_per_s + +# ! ---- Read DICOM images + + +def read_dicom_images() -> darsia.Image: + """Read space-time image from DICOM format for fractip a.""" + + # Provide folder with dicom images (note: two folders) + root = Path( + "/home/jakub/images/ift/fracflow/01-single-phase-tracer-in-fracture/fractip-b" + ) + images = [ + root / Path(p) / Path("DICOM/PA1/ST1/SE1") + for p in [ + "fractip b closed 1 min rekon start 0 15 frames", + ] + ] + + if False: + # Stack images (need to use such approach due to lack of common reference + dicom_images = [] + for img in images: + dicom_image = darsia.imread(img, suffix=".dcm", dim=3, series=True) + dicom_images.append(dicom_image) + uncorrected_image_3d = darsia.stack(dicom_images) + + # Use an assistant to tune the rotation correction and subregion selection + test_image = uncorrected_image_3d.time_slice(9) + test_image.show( + surpress_3d=True, mode="matplotlib", side_view="voxel", cmap="turbo" + ) + test_image.show( + mode="plotly", + side_view="voxel", + view="voxel", + cmap="turbo", + threshold=0.8, + relative=True, + ) + options = { + "threshold": 0.05, + "relative": True, + "verbosity": True, + } + rotation_assistant = darsia.RotationCorrectionAssistant(test_image, **options) + rotation_corrections = rotation_assistant() + rotated_test_image = test_image.copy() + for rotation in rotation_corrections: + rotated_test_image = rotation(rotated_test_image) + subregion_assistant = darsia.SubregionAssistant(rotated_test_image, **options) + coordinates = subregion_assistant() + + else: + # Alternative: Fix tailored input parameters for the correction objects. + rotation_corrections = [ + darsia.RotationCorrection( + **{ + "anchor": np.array([91, 106, 13]), + "rotation_from_isometry": True, + "pts_src": np.array([[91, 106, 13], [97, 106, 86]]), + "pts_dst": np.array([[91, 106, 13], [91, 106, 86]]), + } + ), + darsia.RotationCorrection( + **{ + "anchor": np.array([50, 96, 117]), + "rotation_from_isometry": True, + "pts_src": np.array([[50, 96, 117], [90, 97, 117]]), + "pts_dst": np.array([[50, 96, 117], [90, 96, 117]]), + } + ), + darsia.RotationCorrection( + **{ + "anchor": np.array([106, 94, 0]), + "rotation_from_isometry": True, + "pts_src": np.array([[106, 94, 0], [106, 96, 70]]), + "pts_dst": np.array([[106, 94, 0], [106, 95, 69]]), + } + ), + ] + coordinates = np.array( + [ + [42.21969627956989, 42.29322659139785, -195.13738119462366], + [42.30791713978495, 42.36046250537635, -195.1190467860215], + ] + ) + + # Re-read the dicom image. + # Use the uncorrected images further. + image_3d = darsia.imread( + images, + suffix=".dcm", + dim=3, + series=True, + ) + image_3d = image_3d.time_slice(9) + for correction in rotation_corrections: + image_3d = correction(image_3d) + image_3d = image_3d.subregion(coordinates=coordinates) + + # Plot side view to check the result + if True: + image_3d.show( + mode="matplotlib", + surpress_3d=True, + side_view="voxel", + threshold=0.05, + relative=True, + ) + + # ! ---- Precondition + image_3d.img /= np.max(image_3d.img) + + return image_3d + + +def plot_3d_image(image_3d): + # Plot side view to check the result + image_3d.show( + mode="matplotlib", + surpress_3d=True, + side_view="voxel", + threshold=0.05, + relative=True, + ) + + +def plot_sum(image_3d): + if not isinstance(image_3d, list): + image_3d = [image_3d] + max_val = max( + [ + np.max( + np.sum(img, axis=0) + if isinstance(img, np.ndarray) + else np.sum(img.img, axis=0) + ) + for img in image_3d + ] + ) + for i, img in enumerate(image_3d): + plt.figure(f"sum {i}") + if isinstance(img, np.ndarray): + plt.imshow(np.sum(img, axis=0), cmap="turbo", vmin=0, vmax=max_val) + else: + plt.imshow(np.sum(img.img, axis=0), cmap="turbo", vmin=0, vmax=max_val) + plt.colorbar() + plt.show() + + +def plot_slice(image_3d): + if not isinstance(image_3d, list): + image_3d = [image_3d] + max_val = max( + [np.max(img if isinstance(img, np.ndarray) else img.img) for img in image_3d] + ) + for i, img in enumerate(image_3d): + plt.figure(f"slice {i}") + if isinstance(img, np.ndarray): + plt.imshow(img[20], cmap="turbo", vmin=0, vmax=max_val) + else: + plt.imshow(img.img[20], cmap="turbo", vmin=0, vmax=max_val) + plt.colorbar() + plt.show() + + +# Enrich the signal under the assumption of a monotone displacement experiment. +def accumulate_signal(image: darsia.Image) -> darsia.Image: + new_image = image.copy() + for time_id in range(image.time_num): + new_image.img[..., time_id] = np.max( + image.img[..., slice(0, time_id + 1)], axis=-1 + ) + + return new_image + + +# ! ---- Apply regularization +def apply_tvd(image: darsia.Image) -> darsia.Image: + """NOTE: Most time-consuming routine.""" + + pass + + +def read_regularized_images(image: darsia.Image) -> darsia.Image: + """Read preprocessed regularized images from numpy format.""" + + reg_image = image.copy() + + for i in range(0, 13): + reg_image.img[..., i] = np.load( + # f"tvd/block-a/heterogeneous_mu_1000_iter/tvd_a_1000_{i}.npy" + # f"tvd/block-a/mu_0.0001_1000_iter/tvd_{i}.npy" # + f"tvd/delta/block-a/mu_0.0001_10000_iter/tvd_{i}.npy" # latest tvd + ) + + return reg_image + + +# ! ---- Reduce 3d image to two spatial dimensions +def dimensional_reduction(image_3d: darsia.Image) -> darsia.Image: + """Flatten 3d image.""" + return darsia.reduce_axis(image_3d, axis="z", mode="average") + + +# ! ---- Rescale images to concentrations + + +def extract_concentration( + image_2d: darsia.Image, calibration_interval, injection_rate: float +) -> darsia.Image: + # Define default concentration analysis, enabling calibration + class CalibratedConcentrationAnalysis( + darsia.ConcentrationAnalysis, darsia.InjectionRateModelObjectiveMixin + ): + pass + + model = darsia.ScalingModel(scaling=1) + dicom_concentration_analysis = CalibratedConcentrationAnalysis( + base=None, model=model + ) + + # Calibrate concentration analysis + shape_meta = image_2d.shape_metadata() + geometry = darsia.ExtrudedGeometry(expansion=depth, **shape_meta) + + dicom_concentration_analysis.calibrate_model( + images=image_2d.time_interval(calibration_interval), + options={ + "geometry": geometry, + "injection_rate": injection_rate, + "initial_guess": [1.0], + "tol": 1e-5, + "maxiter": 100, + "method": "Nelder-Mead", + "regression_type": "ransac", + }, + plot_result=True, + ) + + # Interpret the calibration results to effectively determine the injection start. + reference_time = dicom_concentration_analysis.model_calibration_postanalysis() + + # The results is that the effective injection start occurs at 3.25 minutes. Thus, update + # the reference time, or simply the relative time. Set reference time with respect to + # the first active image. Use seconds. + image_2d.update_reference_time(reference_time) + + # Print some of the specs + if True: + print( + f"The identified reference time / injection start is: {reference_time} [s]" + ) + print( + f"The dimensions of the space time dicom image are: {image_2d.dimensions}" + ) + print(f"Relative times in seconds: {image_2d.time}") + + # Only for plotting reasons - same as above but with updated times and with activated plot: + if False: + dicom_concentration_analysis.calibrate_model( + images=image_2d.time_interval( + calibration_interval + ), # Only use some of the first images due to cut-off in the geometry + options={ + "geometry": geometry, + "injection_rate": injection_rate, + "initial_guess": [1.0], + "tol": 1e-5, + "maxiter": 100, + "method": "Nelder-Mead", + "regression_type": "ransac", + }, + plot_result=True, + ) + + # Produce space time image respresentation of concentrations. + dicom_concentration = dicom_concentration_analysis(image_2d) + + # Debug: Plot the concentration profile. Note that values are larger than 1, which is + # due to the calibration based on a non-smooth signal. The expected volume has to be + # distributed among the active pixels. Pixels with strong activity thereby hold + # unphysical amounts of fluid. + if False: + dicom_concentration.show("dicom concentration", 5) + + return dicom_concentration + + +# Determine total concentrations over time +def plot_concentration_evolution(concentration_2d): + shape_meta = concentration_2d.shape_metadata() + geometry = darsia.ExtrudedGeometry(expansion=depth, **shape_meta) + concentration_values = geometry.integrate(concentration_2d) + plt.figure("Experiments - injected volume over time") + plt.plot(concentration_2d.time, concentration_values) + plt.plot( + concentration_2d.time, + [injection_rate * time for time in concentration_2d.time], + color="black", + linestyle=(0, (5, 5)), + ) + plt.xlabel("time [s]") + plt.ylabel("volume [m**3]") + plt.show() + + +# ! ---- VTU images + + +def read_vtu_images(vtu_ind: int) -> darsia.Image: + """Read-in all available VTU data.""" + + # This corresponds approx. to (only temporary - not used further): + vtu_time = 0 # not relevant here. + + # Find the corresponding files + vtu_images_2d = Path(f"data/vtu/block-b/data_2_{str(vtu_ind).zfill(6)}.vtu") + vtu_images_1d = Path(f"data/vtu/block-b/data_1_{str(vtu_ind).zfill(6)}.vtu") + + # Define vtu images as DarSIA images + vtu_image_2d = darsia.imread( + vtu_images_2d, + time=vtu_time, # relative time in minutes + key="temperature", # key to address concentration data + shape=(400, 400), # dim = 2 + vtu_dim=2, # effective dimension of vtu data + ) + + vtu_image_1d = darsia.imread( + vtu_images_1d, + time=vtu_time, # relative time in minutes + key="temperature", # key to address concentration data + shape=(1001, 51), # dim = 2 + vtu_dim=1, # effective dimension of vtu data + width=fracture_aperture, # for embedding in 2d + ) + + # PorePy/meshio cuts off coordinate values at some point... + # Correct manually - width / 2 - aperature / 2. + vtu_image_1d.origin[0] = (6.98 / 2 - 0.1 / 2) * cm2m + + # Equidimensional reconstructions. Superpose 2d and 1d images. + # And directly interpret + porosity_1d = 1.0 - porosity_2d # for volume conservation + vtu_image = darsia.superpose( + [ + darsia.weight(vtu_image_2d, porosity_2d), + darsia.weight(vtu_image_1d, porosity_1d), + ] + ) + + # Plot for testing purposes + if False: + vtu_image.show("equi-dimensionsional reconstruction") + + # Concentrations - simple since the equi-dimensional vtu images are interpreted as + # volumetric concentration + vtu_concentration = vtu_image.copy() + + return vtu_concentration + + +# ! --- Align DICOM and vtu images + + +def align_images(dicom_concentration, vtu_concentration): + # Plot two exemplary images to identify suitable src and dst points which will define + # the alignment procedure + + if True: + plt.figure("dicom") + plt.imshow(np.sum(dicom_concentration.img, axis=0)) + plt.figure("vtu") + plt.imshow(np.sum(vtu_concentration.img, axis=0)) + plt.show() + + # Pixels of fracture end points + voxels_src = [[0, 85, 0], [0, 85, 86], [17, 85, 0], [17, 85, 86]] # in DICOM image + voxels_dst = [ + [0, 237, 86.5], + [0, 148, 86.5], + [17, 237, 86.5], + [17, 148, 86.5], + ] # in VTU image + + # Define coordinate transform and apply it + coordinatesystem_src = dicom_concentration.coordinatesystem + coordinatesystem_dst = vtu_concentration.coordinatesystem + coordinate_transformation = darsia.CoordinateTransformation( + coordinatesystem_src, + coordinatesystem_dst, + voxels_src, + voxels_dst, + fit_options={ + "tol": 1e-5, + "preconditioning": True, + }, + isometry=False, + ) + transformed_dicom_concentration = coordinate_transformation(dicom_concentration) + + # Restrict to intersecting active canvas + intersection = coordinate_transformation.find_intersection() + aligned_dicom_concentration = transformed_dicom_concentration.subregion( + voxels=intersection + ) + aligned_vtu_concentration = vtu_concentration.subregion(voxels=intersection) + + return aligned_dicom_concentration, aligned_vtu_concentration + + +def single_plot( + mode: str, + full_image: darsia.Image, + add_on: str, + image_path: Path, +): + ############################################################################## + # Plot the reconstructed vtu data, vtu plus Gaussian noise, and the dicom data. + + if mode == "sum": + image = darsia.reduce_axis(full_image, axis="z", mode="average") + else: + image = darsia.reduce_axis(full_image, axis="z", mode="slice", slice_idx=20) + + # Define some plotting options (font, x and ylabels) + # matplotlib.rcParams.update({"font.size": 14}) + # Define x and y labels in cm (have to convert from m to cm) + shape = image.num_voxels + dimensions = image.dimensions + + x_pixel, y_pixel = np.meshgrid( + np.linspace(dimensions[0], 0, shape[0]), + np.linspace(0, dimensions[1], shape[1]), + indexing="ij", + ) + vmax = 1.25 + vmin = 0 + cmap = "turbo" + + # Plot the data with contour line. + fig_combination, axs_combination = plt.subplots(nrows=1, ncols=1) + axs_combination.pcolormesh( + y_pixel, x_pixel, image.img, cmap=cmap, vmin=vmin, vmax=vmax + ) + axs_combination.text( + 0.005, 0.075, "experiment", color="white", alpha=0.5, rotation=0, fontsize=14 + ) + axs_combination.text( + 0.005, 0.070, add_on, color="white", alpha=0.5, rotation=0, fontsize=14 + ) + axs_combination.set_ylim(top=0.08) + axs_combination.set_aspect("equal") + axs_combination.set_xlabel("x [m]", fontsize=14) + axs_combination.set_ylabel("y [m]", fontsize=14) + fig_combination.colorbar( + cm.ScalarMappable(cmap=cmap), + ax=axs_combination, + label="volumetric concentration", + ) + fig_combination.savefig(image_path, dpi=500, transparent=True) + + plt.show() + + +def qualitative_comparison( + mode: str, + full_dicom_image: darsia.Image, + full_vtu_image: darsia.Image, + add_on: str, + image_path: Path, + colorbar: bool = False, +): + ############################################################################## + # Plot the reconstructed vtu data, vtu plus Gaussian noise, and the dicom data. + + if mode == "sum": + dicom_image = darsia.reduce_axis(full_dicom_image, axis="z", mode="average") + vtu_image = darsia.reduce_axis(full_vtu_image, axis="z", mode="average") + else: + dicom_image = darsia.reduce_axis( + full_dicom_image, axis="z", mode="slice", slice_idx=20 + ) + vtu_image = darsia.reduce_axis( + full_vtu_image, axis="z", mode="slice", slice_idx=20 + ) + + # Define some plotting options (font, x and ylabels) + # matplotlib.rcParams.update({"font.size": 14}) + # Define x and y labels in cm (have to convert from m to cm) + shape = dicom_image.num_voxels + dimensions = dicom_image.dimensions + + x_pixel, y_pixel = np.meshgrid( + np.linspace(dimensions[0], 0, shape[0]), + np.linspace(0, dimensions[1], shape[1]), + indexing="ij", + ) + vmax = 1.25 + vmin = 0 + contourlevels = [0.045 * vmax, 0.055 * vmax] + cmap = "turbo" + + # Plot the dicom data with contour line. + fig_combination, axs_combination = plt.subplots(nrows=1, ncols=1) + combined_image = dicom_image.img.copy() + mid = 86 + combined_image[:, mid:] = vtu_image.img[:, mid:] + axs_combination.pcolormesh( + y_pixel, x_pixel, combined_image, cmap=cmap, vmin=vmin, vmax=vmax + ) + axs_combination.contourf( + y_pixel, x_pixel, vtu_image.img, cmap="Reds", levels=contourlevels, alpha=0.5 + ) + axs_combination.plot( + [3.45 * cm2m, 3.45 * cm2m], + [0.0, 0.08], + color="white", + alpha=0.3, + linestyle="dashed", + ) + axs_combination.text( + 0.005, 0.075, "experiment", color="white", alpha=0.5, rotation=0, fontsize=14 + ) + axs_combination.text( + 0.005, 0.070, add_on, color="white", alpha=0.5, rotation=0, fontsize=14 + ) + axs_combination.text( + 0.04, 0.075, "simulation", color="white", alpha=0.5, rotation=0, fontsize=14 + ) + axs_combination.set_ylim(top=0.08) + axs_combination.set_aspect("equal") + axs_combination.set_xlabel("x [m]", fontsize=14) + axs_combination.set_ylabel("y [m]", fontsize=14) + if colorbar: + fig_combination.colorbar( + cm.ScalarMappable(cmap=cmap), + ax=axs_combination, + label="volumetric concentration", + ) + fig_combination.savefig(image_path, dpi=500, transparent=True) + + plt.show() + + +########################################################################### +# Main analysis + +# Original PET images +if True: + dicom_image_3d = read_dicom_images() + dicom_image_3d.save("data/npz/block-b/dicom_raw_3d.npz") +else: + dicom_image_3d = darsia.imread("data/npz/block-b/dicom_raw_3d.npz") + +# Pick corresponding vtu images. +vtu_ind = 439 +vtu_2d_concentration = read_vtu_images(vtu_ind) + +# Resize to similar shape as dicom image +dicom_voxel_size = dicom_image_3d.voxel_size +vtu_2d_concentration = darsia.equalize_voxel_size( + vtu_2d_concentration, min(dicom_voxel_size) +) + +# Expand vtu image to 3d +dicom_height = dicom_image_3d.dimensions[0] +dicom_shape = dicom_image_3d.img.shape +vtu_concentration_3d = darsia.extrude_along_axis( + vtu_2d_concentration, dicom_height, dicom_shape[0] +) + +# Align dicom and vtu +if True: + dicom_concentration = dicom_image_3d.copy() + aligned_dicom_concentration, aligned_vtu_concentration = align_images( + dicom_concentration, vtu_concentration_3d + ) + aligned_dicom_concentration.save("data/npz/block-b/aligned_dicom_concentration.npz") + aligned_vtu_concentration.save( + f"data/npz/block-b/aligned_vtu_{vtu_ind}_concentration.npz" + ) +else: + aligned_dicom_concentration = darsia.imread( + "data/npz/block-b/aligned_dicom_concentration.npz" + ) + aligned_vtu_concentration = darsia.imread( + f"data/npz/block-b/aligned_vtu_{vtu_ind}_concentration.npz" + ) + +# Define final vtu concentration, and compute its mass (reference) +vtu_concentration_3d = aligned_vtu_concentration.copy() +vtu_3d_shape = vtu_concentration_3d.shape_metadata() +vtu_3d_geometry = darsia.Geometry(**vtu_3d_shape) +vtu_3d_integral = vtu_3d_geometry.integrate(vtu_concentration_3d) + + +# DICOM without TVD +def rescale_data(image, ref_integral): + shape = image.shape_metadata() + geometry = darsia.Geometry(**shape) + integral = geometry.integrate(image) + image.img *= ref_integral / integral + return image + + +# Define dicom concentration with same mass +dicom_concentration_3d = aligned_dicom_concentration.copy() +dicom_concentration_3d = rescale_data(dicom_concentration_3d, vtu_3d_integral) +# single_plot("slice", dicom_concentration_3d, "noisy", "plots/block-b/lab_data_slice.png") +# single_plot("sum", dicom_concentration_3d, "noisy", "plots/block-b/lab_data_avg.png") + +# Define mask (omega) for trust for regularization +heterogeneous_omega = True +if heterogeneous_omega: + dicom_rescaled = dicom_concentration_3d.copy() + dicom_rescaled.img /= np.max(dicom_rescaled.img) + omega_bound = 0.15 + omega = np.minimum(dicom_rescaled.img, omega_bound) + mask_zero = dicom_rescaled.img < 1e-4 + omega[mask_zero] = 10 + # plot_slice(omega) + + # Save plot + reduced_dicom_concentration_2d = darsia.reduce_axis( + dicom_concentration_3d, axis="z", mode="slice", slice_idx=20 + ) + shape = reduced_dicom_concentration_2d.num_voxels + dimensions = reduced_dicom_concentration_2d.dimensions + + x_pixel, y_pixel = np.meshgrid( + np.linspace(dimensions[0], 0, shape[0]), + np.linspace(0, dimensions[1], shape[1]), + indexing="ij", + ) + fig_mask, axs_mask = plt.subplots(nrows=1, ncols=1) + axs_mask.pcolormesh(y_pixel, x_pixel, np.log(omega[20]), cmap="turbo") + axs_mask.set_ylim(top=0.08) + axs_mask.set_aspect("equal") + axs_mask.set_xlabel("x [m]") # , fontsize=14) + axs_mask.set_ylabel("y [m]") # , fontsize=14) + fig_mask.colorbar( + cm.ScalarMappable(cmap="turbo"), ax=axs_mask, label="log($\omega_f$)" + ) + fig_mask.savefig("plots/block-b/omega.png", dpi=500, transparent=True) + plt.close() +else: + omega = 0.015 + +# Choose file suffix depeneding on omega +suffix = "heter" if heterogeneous_omega else "hom" + +# DICOM concentration with H1 regularization +if True: + h1_reg_dicom_concentration_3d = darsia.H1_regularization( + dicom_concentration_3d, + mu=0.1, + omega=omega, + dim=3, + solver=darsia.CG(maxiter=10000, tol=1e-5), + ) + h1_reg_dicom_concentration_3d.save( + f"data/npz/block-b/h1_reg_dicom_concentration_{suffix}.npz" + ) +else: + h1_reg_dicom_concentration_3d = darsia.imread( + f"data/npz/block-b/h1_reg_dicom_concentration_{suffix}.npz" + ) +h1_reg_dicom_concentration_3d = rescale_data( + h1_reg_dicom_concentration_3d, vtu_3d_integral +) + +# DICOM concentration with TVD regularization +isotropic = True +isotropic_suffix = "iso" if isotropic else "aniso" +if True: + tvd_reg_dicom_concentration_3d = darsia.tvd( + dicom_concentration_3d, + method="heterogeneous bregman", + isotropic=isotropic, + weight=0.02, + omega=omega, + dim=3, + max_num_iter=100, + eps=1e-5, + verbose=True, + solver=darsia.Jacobi(maxiter=20), + ) + tvd_reg_dicom_concentration_3d.save( + f"data/npz/block-b/tvd_{isotropic_suffix}_reg_dicom_concentration_{suffix}.npz" + ) +else: + tvd_reg_dicom_concentration_3d = darsia.imread( + f"data/npz/block-b/tvd_{isotropic_suffix}_reg_dicom_concentration_{suffix}.npz" + ) +tvd_reg_dicom_concentration_3d = rescale_data( + tvd_reg_dicom_concentration_3d, vtu_3d_integral +) + +# Make qualitative comparisons +if True: + # plot_sum( + # [ + # vtu_concentration_3d, + # dicom_concentration_3d, + # h1_reg_dicom_concentration_3d, + # tvd_reg_dicom_concentration_3d, + # ] + # ) + # plot_slice( + # [ + # vtu_concentration_3d, + # dicom_concentration_3d, + # h1_reg_dicom_concentration_3d, + # tvd_reg_dicom_concentration_3d, + # ] + # ) + # qualitative_comparison( + # "sum", + # dicom_concentration_3d, + # vtu_concentration_3d, + # "noisy", + # "plots/block-b/pure_dicom_avg.png", + # ) + qualitative_comparison( + "slice", + dicom_concentration_3d, + vtu_concentration_3d, + "noisy", + "plots/block-b/pure_dicom_slice.png", + ) + # qualitative_comparison( + # "sum", + # h1_reg_dicom_concentration_3d, + # vtu_concentration_3d, + # "H1", + # f"plots/block-b/h1_reg_dicom_avg_{suffix}.png", + # ) + qualitative_comparison( + "slice", + h1_reg_dicom_concentration_3d, + vtu_concentration_3d, + "H1", + f"plots/block-b/h1_reg_dicom_slice_{suffix}.png", + ) + qualitative_comparison( + "sum", + tvd_reg_dicom_concentration_3d, + vtu_concentration_3d, + "TVD", + f"plots/block-b/tvd_{isotropic_suffix}_reg_dicom_avg_{suffix}.png", + colorbar=True, + ) + qualitative_comparison( + "slice", + tvd_reg_dicom_concentration_3d, + vtu_concentration_3d, + "TVD", + f"plots/block-b/tvd_{isotropic_suffix}_reg_dicom_slice_{suffix}.png", + ) + +# Quantitative comparison - Wasserstein distance in 2d (only for illustration purposes) + +slice_idx = 20 +dicom_slice = darsia.reduce_axis( + dicom_concentration_3d, axis="z", mode="slice", slice_idx=slice_idx +) +vtu_slice = darsia.reduce_axis( + vtu_concentration_3d, axis="z", mode="slice", slice_idx=slice_idx +) +h1_reg_dicom_slice = darsia.reduce_axis( + h1_reg_dicom_concentration_3d, axis="z", mode="slice", slice_idx=slice_idx +) +tvd_reg_dicom_slice = darsia.reduce_axis( + tvd_reg_dicom_concentration_3d, axis="z", mode="slice", slice_idx=slice_idx +) + + +# Rescale (unphysical but necessary for comparison) +def rescale_slice(image: darsia.Image, ref_integral) -> darsia.Image: + shape = image.shape_metadata() + geometry = darsia.Geometry(**shape) + integral = geometry.integrate(image) + image.img *= ref_integral / integral + return image + + +# Determine reference value +shape = vtu_slice.shape_metadata() +geometry = darsia.Geometry(**shape) +ref_integral = geometry.integrate(vtu_slice) + +dicom_slice = rescale_slice(dicom_slice, ref_integral) +h1_reg_dicom_slice = rescale_slice(h1_reg_dicom_slice, ref_integral) +tvd_reg_dicom_slice = rescale_slice(tvd_reg_dicom_slice, ref_integral) + +# Wasserstein for 2d images +options = { + "method": "newton", + "L": 1e-2, + "num_iter": 500, + "tol": 1e-11, + "tol_distance": 1e-12, + "regularization": 1e-16, + "depth": 10, + "lumping": True, + "verbose": True, + "scaling": 10, + "plot_frequency": 10, +} +kwargs = { + "method": "newton", + "options": options, + "plot_solution": True, +} +distance_dicom_vtu = darsia.wasserstein_distance( + dicom_slice, vtu_slice, name="noisy vs. simulation", **kwargs +) +distance_h1_dicom_vtu = darsia.wasserstein_distance( + h1_reg_dicom_slice, vtu_slice, name="H1 vs simulation", **kwargs +) +distance_tvd_dicom_vtu = darsia.wasserstein_distance( + tvd_reg_dicom_slice, vtu_slice, name="TVD vs simulation", **kwargs +) + +print("The distances:") +print("Pure DICOM: ", distance_dicom_vtu) +print("H1 reg DICOM: ", distance_h1_dicom_vtu) +print("TVD reg DICOM: ", distance_tvd_dicom_vtu) + +assert False, "3d Wasserstein distance computations not sufficiently efficient." + +# Quantitative comparison - Wasserstein distance in 3d +options = { + "method": "newton", + "L": 1e-2, + "num_iter": 100, + "tol": 1e-10, + "tol_distance": 1e-11, + "regularization": 1e-16, + "depth": 10, + "lumping": True, + "verbose": True, +} +kwargs = { + "method": "newton", + "options": options, + "plot_solution": False, +} +distance_dicom_vtu_3d = darsia.wasserstein_distance_3d( + dicom_concentration_3d, vtu_concentration_3d, name="", **kwargs +) +distance_h1_dicom_vtu_3d = darsia.wasserstein_distance_3d( + h1_reg_dicom_concentration_3d, + vtu_concentration_3d, + name="", + **kwargs, +) +distance_tvd_dicom_vtu_3d = darsia.wasserstein_distance_3d( + tvd_reg_dicom_concentration_3d, + vtu_concentration_3d, + name="", + **kwargs, +) + +print("The distances 3d:") +print("Pure DICOM: ", distance_dicom_vtu_3d) +print("H1 reg DICOM: ", distance_h1_dicom_vtu_3d) +print("TVD reg DICOM: ", distance_tvd_dicom_vtu_3d) diff --git a/requirements-dev.txt b/requirements-dev.txt index 63eef17d..4824de80 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,3 +5,4 @@ flake8 mypy meshio deepdiff +pyamg diff --git a/src/darsia/__init__.py b/src/darsia/__init__.py index 122e1bcb..e6ba1a2f 100644 --- a/src/darsia/__init__.py +++ b/src/darsia/__init__.py @@ -25,6 +25,8 @@ from darsia.utils.linear_solvers.mg import * from darsia.utils.andersonacceleration import * from darsia.utils.dtype import * +from darsia.utils.grid import * +from darsia.utils.fv import * from darsia.corrections.basecorrection import * from darsia.corrections.shape.curvature import * from darsia.corrections.shape.affine import * @@ -67,3 +69,5 @@ from darsia.assistants.base_assistant import * from darsia.assistants.rotation_correction_assistant import * from darsia.assistants.subregion_assistant import * + +# from darsia.assistants.curvature_correction_assistant import * diff --git a/src/darsia/image/arithmetics.py b/src/darsia/image/arithmetics.py index a2a38aa1..d2bc199f 100644 --- a/src/darsia/image/arithmetics.py +++ b/src/darsia/image/arithmetics.py @@ -32,9 +32,10 @@ def weight(img: darsia.Image, weight: Union[float, int, darsia.Image]) -> darsia weighted_img.img *= weight elif isinstance(weight, darsia.Image): - assert darsia.check_equal_coordinatesystems( + equal_coordinate_system, log = darsia.check_equal_coordinatesystems( img.coordinatesystem, weight.coordinatesystem, exclude_size=True ) + assert equal_coordinate_system, f"{log}" space_dim = img.space_dim assert len(weight.img.shape) == space_dim diff --git a/src/darsia/image/coordinatesystem.py b/src/darsia/image/coordinatesystem.py index bc5cd830..75e94428 100644 --- a/src/darsia/image/coordinatesystem.py +++ b/src/darsia/image/coordinatesystem.py @@ -225,7 +225,7 @@ def check_equal_coordinatesystems( coordinatesystem1: CoordinateSystem, coordinatesystem2: CoordinateSystem, exclude_size: bool = False, -) -> bool: +) -> tuple[bool, dict]: """Check whether two coordinate systems are equivalent, i.e., they share basic attributes. @@ -236,38 +236,64 @@ def check_equal_coordinatesystems( Returns: bool: True iff the two coordinate systems are equivalent. + dict: log of the failed checks. """ success = True - success = success and coordinatesystem1.indexing == coordinatesystem2.indexing - success = success and coordinatesystem1.dim == coordinatesystem2.dim + failure_log = [] + + if not (coordinatesystem1.indexing == coordinatesystem2.indexing): + failure_log.append("indexing") + + if not (coordinatesystem1.dim == coordinatesystem2.dim): + failure_log.append("dim") + if not exclude_size: - success = success and np.allclose( - coordinatesystem1.shape, coordinatesystem2.shape - ) - success = success and np.allclose( - coordinatesystem1.dimensions, coordinatesystem2.dimensions - ) - success = success and coordinatesystem1.axes == coordinatesystem2.axes + if not (np.allclose(coordinatesystem1.shape, coordinatesystem2.shape)): + failure_log.append("shape") + + if not (np.allclose(coordinatesystem1.dimensions, coordinatesystem2.dimensions)): + failure_log.append("dimensions") + + if not (coordinatesystem1.axes == coordinatesystem2.axes): + failure_log.append("axes") + if not exclude_size: + voxel_size_equal = True for axis in coordinatesystem1.axes: - success = ( - success + voxel_size_equal = ( + voxel_size_equal and coordinatesystem1.voxel_size[axis] == coordinatesystem2.voxel_size[axis] ) - success = success and np.allclose( - coordinatesystem1._coordinate_of_origin_voxel, - coordinatesystem2._coordinate_of_origin_voxel, - ) - success = success and np.allclose( - coordinatesystem1._coordinate_of_opposite_voxel, - coordinatesystem2._coordinate_of_opposite_voxel, - ) - if not exclude_size: - success = success and np.allclose( - coordinatesystem1._voxel_of_origin_coordinate, - coordinatesystem2._voxel_of_origin_coordinate, + if not voxel_size_equal: + failure_log.append("voxel_size") + + if not ( + np.allclose( + coordinatesystem1._coordinate_of_origin_voxel, + coordinatesystem2._coordinate_of_origin_voxel, + ) + ): + failure_log.append("coordinate_of_origin_voxel") + + if not ( + np.allclose( + coordinatesystem1._coordinate_of_opposite_voxel, + coordinatesystem2._coordinate_of_opposite_voxel, ) + ): + failure_log.append("coordinate_of_opposite_voxel") + + if not exclude_size: + if not ( + np.allclose( + coordinatesystem1._voxel_of_origin_coordinate, + coordinatesystem2._voxel_of_origin_coordinate, + ) + ): + failure_log.append("voxel_of_origin_coordinate") + + success = len(failure_log) == 0 - return success + return success, failure_log diff --git a/src/darsia/measure/emd.py b/src/darsia/measure/emd.py index 69987c46..e751f984 100644 --- a/src/darsia/measure/emd.py +++ b/src/darsia/measure/emd.py @@ -46,6 +46,10 @@ def __call__( float or array: distance between img_1 and img_2. """ + # Two-dimensional + if not (img_1.space_dim == 2 and img_2.space_dim == 2): + raise NotImplementedError("EMD only implemented for 2d.") + # FIXME investigation required regarding resize preprocessing... # Preprocess images preprocessed_img_1 = self._preprocess(img_1) @@ -116,13 +120,11 @@ def _compatibility_check( # Series assert img_1.time_num == img_2.time_num - # Two-dimensional - assert img_1.space_dim == 2 and img_2.space_dim == 2 - # Check whether the coordinate system is compatible - assert darsia.check_equal_coordinatesystems( + equal_coordinate_system, log = darsia.check_equal_coordinatesystems( img_1.coordinatesystem, img_2.coordinatesystem ) + assert equal_coordinate_system, f"{log}" assert np.allclose(img_1.voxel_size, img_2.voxel_size) # Compatible distributions - comparing sums is sufficient since it is implicitly diff --git a/src/darsia/measure/wasserstein.py b/src/darsia/measure/wasserstein.py index 655db865..2e1f4ac3 100644 --- a/src/darsia/measure/wasserstein.py +++ b/src/darsia/measure/wasserstein.py @@ -4,13 +4,24 @@ from __future__ import annotations import time +from pathlib import Path +from typing import Optional, Union import matplotlib.pyplot as plt import numpy as np +import pyamg import scipy.sparse as sps import darsia +# General TODO list +# - improve documentation, in particular with focus on keywords +# - remove plotting +# - improve assembling of operators through partial assembling +# - improve stopping criteria +# - use better quadrature for l1_dissipation? +# - allow to reuse setup. + class VariationalWassersteinDistance(darsia.EMD): """Base class for setting up the variational Wasserstein distance. @@ -30,389 +41,659 @@ class VariationalWassersteinDistance(darsia.EMD): def __init__( self, - shape: tuple, - voxel_size: list, - dim: int, + grid: darsia.Grid, options: dict = {}, ) -> None: """ Args: - shape (tuple): shape of the image - voxel_size (list): voxel size of the image - dim (int): dimension of the problem + grid (darsia.Grid): tensor grid associated with the images options (dict): options for the solver - num_iter (int): maximum number of iterations. Defaults to 100. - tol (float): tolerance for the stopping criterion. Defaults to 1e-6. - L (float): parameter for the Bregman iteration. Defaults to 1.0. - regularization (float): regularization parameter for the Bregman iteration. Defaults to 0.0. - - depth (int): depth of the Anderson acceleration. Defaults to 0. + - aa_depth (int): depth of the Anderson acceleration. Defaults to 0. + - aa_restart (int): restart of the Anderson acceleration. Defaults to None. - scaling (float): scaling of the fluxes in the plot. Defaults to 1.0. - lumping (bool): lump the mass matrix. Defaults to True. """ - # TODO improve documentation for options - method dependent # Cache geometrical infos - self.shape = shape - self.voxel_size = voxel_size - self.dim = dim + self.grid = grid + """darsia.Grid: grid""" - assert dim == 2, "Currently only 2D images are supported." + self.voxel_size = grid.voxel_size + """np.ndarray: voxel size""" self.options = options + """dict: options for the solver""" + self.regularization = self.options.get("regularization", 0.0) + """float: regularization parameter""" + self.verbose = self.options.get("verbose", False) + """bool: verbosity""" - # Setup - self._setup() - - def _setup(self) -> None: - """Setup of fixed discretization""" - - # Define dimensions of the problem - dim_cells = self.shape - num_cells = np.prod(dim_cells) - numbering_cells = np.arange(num_cells, dtype=int).reshape(dim_cells) - - # Consider only inner edges - vertical_edges_shape = (self.shape[0], self.shape[1] - 1) - horizontal_edges_shape = (self.shape[0] - 1, self.shape[1]) - num_edges_axis = [ - np.prod(vertical_edges_shape), - np.prod(horizontal_edges_shape), - ] - num_edges = np.sum(num_edges_axis) - - # Define connectivity - connectivity = np.zeros((num_edges, 2), dtype=int) - connectivity[: num_edges_axis[0], 0] = np.ravel( - numbering_cells[:, :-1] - ) # left cells - connectivity[: num_edges_axis[0], 1] = np.ravel( - numbering_cells[:, 1:] - ) # right cells - connectivity[num_edges_axis[0] :, 0] = np.ravel( - numbering_cells[:-1, :] - ) # top cells - connectivity[num_edges_axis[0] :, 1] = np.ravel( - numbering_cells[1:, :] - ) # bottom cells - - # Define sparse divergence operator, integrated over elements: flat_fluxes -> flat_mass - div_data = np.concatenate( - ( - self.voxel_size[0] * np.ones(num_edges_axis[0], dtype=float), - self.voxel_size[1] * np.ones(num_edges_axis[1], dtype=float), - -self.voxel_size[0] * np.ones(num_edges_axis[0], dtype=float), - -self.voxel_size[1] * np.ones(num_edges_axis[1], dtype=float), - ) - ) - div_row = np.concatenate( - ( - connectivity[: num_edges_axis[0], 0], - connectivity[num_edges_axis[0] :, 0], - connectivity[: num_edges_axis[0], 1], - connectivity[num_edges_axis[0] :, 1], - ) - ) - div_col = np.tile(np.arange(num_edges, dtype=int), 2) - self.div = sps.csc_matrix( - (div_data, (div_row, div_col)), shape=(num_cells, num_edges) + # Setup of method + self._setup_dof_management() + self._setup_discretization() + self._setup_linear_solver() + self._setup_acceleration() + + def _setup_dof_management(self) -> None: + """Setup of DOF management. + + The following degrees of freedom are considered (also in this order): + - flat fluxes (normal fluxes on the faces) + - flat potentials (potentials on the cells) + - lagrange multiplier (scalar variable) - Idea: Fix the potential in the + center of the domain to zero. This is done by adding a constraint to the + potential via a Lagrange multiplier. + + """ + # ! ---- Number of dofs ---- + num_flux_dofs = self.grid.num_faces + num_potential_dofs = self.grid.num_cells + num_lagrange_multiplier_dofs = 1 + num_dofs = num_flux_dofs + num_potential_dofs + num_lagrange_multiplier_dofs + + # ! ---- Indices in global system ---- + self.flux_indices = np.arange(num_flux_dofs) + """np.ndarray: indices of the fluxes""" + + self.potential_indices = np.arange( + num_flux_dofs, num_flux_dofs + num_potential_dofs ) + """np.ndarray: indices of the potentials""" - # Define sparse mass matrix on cells: flat_mass -> flat_mass - self.mass_matrix_cells = sps.diags( - np.prod(self.voxel_size) * np.ones(num_cells, dtype=float) + self.lagrange_multiplier_indices = np.array( + [num_flux_dofs + num_potential_dofs], dtype=int ) + """np.ndarray: indices of the lagrange multiplier""" - # Define sparse mass matrix on edges: flat_fluxes -> flat_fluxes - lumping = self.options.get("lumping", True) - if lumping: - self.mass_matrix_edges = sps.diags( - np.prod(self.voxel_size) * np.ones(num_edges, dtype=float) - ) - else: - # Define connectivity: cell to face (only for inner cells) - connectivity_cell_to_vertical_face = np.zeros((num_cells, 2), dtype=int) - connectivity_cell_to_vertical_face[ - np.ravel(numbering_cells[:, :-1]), 0 - ] = np.arange( - num_edges_axis[0] - ) # left face - connectivity_cell_to_vertical_face[ - np.ravel(numbering_cells[:, 1:]), 1 - ] = np.arange( - num_edges_axis[0] - ) # right face - connectivity_cell_to_horizontal_face = np.zeros((num_cells, 2), dtype=int) - connectivity_cell_to_horizontal_face[ - np.ravel(numbering_cells[:-1, :]), 0 - ] = np.arange( - num_edges_axis[0], num_edges_axis[0] + num_edges_axis[1] - ) # top face - connectivity_cell_to_horizontal_face[ - np.ravel(numbering_cells[1:, :]), 1 - ] = np.arange( - num_edges_axis[0], num_edges_axis[0] + num_edges_axis[1] - ) # bottom face - - # Info about inner cells - inner_cells_with_vertical_faces = np.ravel(numbering_cells[:, 1:-1]) - inner_cells_with_horizontal_faces = np.ravel(numbering_cells[1:-1, :]) - num_inner_cells_with_vertical_faces = len(inner_cells_with_vertical_faces) - num_inner_cells_with_horizontal_faces = len( - inner_cells_with_horizontal_faces - ) + # ! ---- Fast access to components through slices ---- + self.flux_slice = slice(0, num_flux_dofs) + """slice: slice for the fluxes""" - # Define true RT0 mass matrix on edges: flat_fluxes -> flat_fluxes - mass_matrix_edges_data = np.prod(self.voxel_size) * np.concatenate( - ( - 2 / 3 * np.ones(num_edges, dtype=float), # all faces - 1 - / 6 - * np.ones( - num_inner_cells_with_vertical_faces, dtype=float - ), # left faces - 1 - / 6 - * np.ones( - num_inner_cells_with_vertical_faces, dtype=float - ), # right faces - 1 - / 6 - * np.ones( - num_inner_cells_with_horizontal_faces, dtype=float - ), # top faces - 1 - / 6 - * np.ones( - num_inner_cells_with_horizontal_faces, dtype=float - ), # bottom faces - ) - ) - mass_matrix_edges_row = np.concatenate( - ( - np.arange(num_edges, dtype=int), - connectivity_cell_to_vertical_face[ - inner_cells_with_vertical_faces, 0 - ], - connectivity_cell_to_vertical_face[ - inner_cells_with_vertical_faces, 1 - ], - connectivity_cell_to_horizontal_face[ - inner_cells_with_horizontal_faces, 0 - ], - connectivity_cell_to_horizontal_face[ - inner_cells_with_horizontal_faces, 1 - ], - ) - ) - mass_matrix_edges_col = np.concatenate( - ( - np.arange(num_edges, dtype=int), - connectivity_cell_to_vertical_face[ - inner_cells_with_vertical_faces, 1 - ], - connectivity_cell_to_vertical_face[ - inner_cells_with_vertical_faces, 0 - ], - connectivity_cell_to_horizontal_face[ - inner_cells_with_horizontal_faces, 1 - ], - connectivity_cell_to_horizontal_face[ - inner_cells_with_horizontal_faces, 0 - ], - ) - ) - self.mass_matrix_edges = sps.csc_matrix( - ( - mass_matrix_edges_data, - (mass_matrix_edges_row, mass_matrix_edges_col), - ), - shape=(num_edges, num_edges), - ) + self.potential_slice = slice(num_flux_dofs, num_flux_dofs + num_potential_dofs) + """slice: slice for the potentials""" - # Utilities - depth = self.options.get("depth", 0) - self.anderson = ( - darsia.AndersonAcceleration(dimension=num_edges, depth=depth) - if depth > 0 - else None + self.lagrange_multiplier_slice = slice( + num_flux_dofs + num_potential_dofs, + num_flux_dofs + num_potential_dofs + num_lagrange_multiplier_dofs, ) + """slice: slice for the lagrange multiplier""" - # TODO needs to be defined for each problem separately + self.reduced_system_slice = slice(num_flux_dofs, None) + """slice: slice for the reduced system (potentials and lagrange multiplier)""" - # Define sparse embedding operator for fluxes into full discrete DOF space + # Embedding operators self.flux_embedding = sps.csc_matrix( ( - np.ones(num_edges, dtype=float), - (np.arange(num_edges), np.arange(num_edges)), + np.ones(num_flux_dofs, dtype=float), + (self.flux_indices, self.flux_indices), ), - shape=(num_edges + num_cells + 1, num_edges), + shape=(num_dofs, num_flux_dofs), ) + """sps.csc_matrix: embedding operator for fluxes""" - # Cache - self.num_edges = num_edges - self.num_cells = num_cells - self.dim_cells = dim_cells - self.numbering_cells = numbering_cells - self.num_edges_axis = num_edges_axis - self.vertical_edges_shape = vertical_edges_shape - self.horizontal_edges_shape = horizontal_edges_shape + def _setup_discretization(self) -> None: + """Setup of fixed discretization operators.""" - def _problem_specific_setup(self, mass_diff: np.ndarray) -> None: - """Resetup of fixed discretization""" + # ! ---- Constraint for the potential correpsonding to Lagrange multiplier ---- - # TODO can't we just fix some cell, e.g., [0,0]. Move then this to the above. + center_cell = np.array(self.grid.shape) // 2 + self.constrained_cell_flat_index = np.ravel_multi_index( + center_cell, self.grid.shape + ) + """int: flat index of the cell where the potential is constrained to zero""" - # Fix index of dominating contribution in image differece - self.constrained_cell_flat_index = np.argmax(np.abs(mass_diff)) - self.pressure_constraint = sps.csc_matrix( + num_potential_dofs = self.grid.num_cells + self.potential_constraint = sps.csc_matrix( ( np.ones(1, dtype=float), (np.zeros(1, dtype=int), np.array([self.constrained_cell_flat_index])), ), - shape=(1, self.num_cells), + shape=(1, num_potential_dofs), dtype=float, ) + """sps.csc_matrix: effective constraint for the potential""" + + # ! ---- Discretization operators ---- + + self.div = darsia.FVDivergence(self.grid).mat + """sps.csc_matrix: divergence operator: flat fluxes -> flat potentials""" + + self.mass_matrix_cells = darsia.FVMass(self.grid).mat + """sps.csc_matrix: mass matrix on cells: flat potentials -> flat potentials""" + + lumping = self.options.get("lumping", True) + self.mass_matrix_faces = darsia.FVMass(self.grid, "faces", lumping).mat + """sps.csc_matrix: mass matrix on faces: flat fluxes -> flat fluxes""" + + self.face_reconstruction = darsia.FVFullFaceReconstruction(self.grid) + """sps.csc_matrix: full face reconstruction: flat fluxes -> vector fluxes""" - # Linear part of the operator. + # Linear part of the Darcy operator with potential constraint. self.broken_darcy = sps.bmat( [ [None, -self.div.T, None], - [self.div, None, -self.pressure_constraint.T], - [None, self.pressure_constraint, None], + [self.div, None, -self.potential_constraint.T], + [None, self.potential_constraint, None], ], format="csc", ) + """sps.csc_matrix: linear part of the Darcy operator""" - def split_solution( - self, solution: np.ndarray - ) -> tuple[np.ndarray, np.ndarray, float]: - """Split the solution into fluxes, pressure and lagrange multiplier. + L_init = self.options.get("L_init", 1.0) + self.darcy_init = sps.bmat( + [ + [L_init * self.mass_matrix_faces, -self.div.T, None], + [self.div, None, -self.potential_constraint.T], + [None, self.potential_constraint, None], + ], + format="csc", + ) + """sps.csc_matrix: initial Darcy operator""" + + def _setup_linear_solver(self) -> None: + self.linear_solver_type = self.options.get("linear_solver", "lu") + assert self.linear_solver_type in [ + "lu", + "lu-flux-reduced", + "amg-flux-reduced", + "lu-potential", + "amg-potential", + ], f"Linear solver {self.linear_solver_type} not supported." + """str: type of linear solver""" + + if self.linear_solver_type in ["amg-flux-reduced", "amg-potential"]: + # TODO add possibility for user control + self.ml_options = { + # B=X.reshape( + # n * n, 1 + # ), # the representation of the near null space (this is a poor choice) + # BH=None, # the representation of the left near null space + "symmetry": "hermitian", # indicate that the matrix is Hermitian + # strength="evolution", # change the strength of connection + "aggregate": "standard", # use a standard aggregation method + "smooth": ( + "jacobi", + {"omega": 4.0 / 3.0, "degree": 2}, + ), # prolongation smoothing + "presmoother": ("block_gauss_seidel", {"sweep": "symmetric"}), + "postsmoother": ("block_gauss_seidel", {"sweep": "symmetric"}), + # improve_candidates=[ + # ("block_gauss_seidel", {"sweep": "symmetric", "iterations": 4}), + # None, + # ], + "max_levels": 4, # maximum number of levels + "max_coarse": 1000, # maximum number on a coarse level + # keep=False, # keep extra operators around in the hierarchy (memory) + } + """dict: options for the AMG solver""" + + self.tol_amg = self.options.get("linear_solver_tol", 1e-6) + """float: tolerance for the AMG solver""" + + self.res_history_amg = [] + """list: history of residuals for the AMG solver""" + + # Setup inrastructure for Schur complement reduction + if self.linear_solver_type in ["lu-flux-reduced", "amg-flux-reduced"]: + self.setup_one_level_schur_reduction() + + elif self.linear_solver_type in ["lu-potential", "amg-potential"]: + self.setup_two_level_schur_reduction() + + def setup_one_level_schur_reduction(self) -> None: + """Setup the infrastructure for reduced systems through Gauss elimination. + + Provide internal data structures for the reduced system. - Args: - solution (np.ndarray): solution + """ + # Step 1: Compute the jacobian of the Darcy problem - Returns: - tuple: fluxes, pressure, lagrange multiplier + jacobian = self.darcy_init.copy() - """ - # Split the solution - flat_flux = solution[: self.num_edges] - flat_pressure = solution[self.num_edges : self.num_edges + self.num_cells] - flat_lagrange_multiplier = solution[-1] + # Step 2: Remove flux blocks through Schur complement approach + + # Build Schur complement wrt. flux-flux block + J_inv = sps.diags(1.0 / jacobian.diagonal()[self.flux_slice]) + D = jacobian[self.reduced_system_slice, self.flux_slice].copy() + schur_complement = D.dot(J_inv.dot(D.T)) + + # Cache divergence matrix + self.D = D.copy() + """sps.csc_matrix: divergence matrix""" + + self.DT = self.D.T.copy() + """sps.csc_matrix: transposed divergence matrix""" + + # Cache (constant) jacobian subblock + self.jacobian_subblock = jacobian[ + self.reduced_system_slice, self.reduced_system_slice + ].copy() + """sps.csc_matrix: constant jacobian subblock of the reduced system""" + + # Add Schur complement - use this to identify sparsity structure + # Cache the reduced jacobian + self.reduced_jacobian = self.jacobian_subblock + schur_complement + """sps.csc_matrix: reduced jacobian incl. Schur complement""" + + def setup_two_level_schur_reduction(self) -> None: + """Additional setup of infrastructure for fully reduced systems.""" + # Step 1 and 2: + self.setup_one_level_schur_reduction() + + # Step 3: Remove Lagrange multiplier block through Gauss elimination + + # Find row entries to be removed + rm_row_entries = np.arange( + self.reduced_jacobian.indptr[self.constrained_cell_flat_index], + self.reduced_jacobian.indptr[self.constrained_cell_flat_index + 1], + ) + + # Find column entries to be removed + rm_col_entries = np.where( + self.reduced_jacobian.indices == self.constrained_cell_flat_index + )[0] + + # Collect all entries to be removes + rm_indices = np.unique( + np.concatenate((rm_row_entries, rm_col_entries)).astype(int) + ) + # Cache for later use in remove_lagrange_multiplier + self.rm_indices = rm_indices + """np.ndarray: indices to be removed in the reduced system""" + + # Identify rows to be reduced + rm_rows = [ + np.max(np.where(self.reduced_jacobian.indptr <= index)[0]) + for index in rm_indices + ] + + # Reduce data - simply remove + fully_reduced_jacobian_data = np.delete(self.reduced_jacobian.data, rm_indices) + + # Reduce indices - remove and shift + fully_reduced_jacobian_indices = np.delete( + self.reduced_jacobian.indices, rm_indices + ) + fully_reduced_jacobian_indices[ + fully_reduced_jacobian_indices > self.constrained_cell_flat_index + ] -= 1 + + # Reduce indptr - shift and remove + # NOTE: As only a few entries should be removed, this is not too expensive + # and a for loop is used + fully_reduced_jacobian_indptr = self.reduced_jacobian.indptr.copy() + for row in rm_rows: + fully_reduced_jacobian_indptr[row + 1 :] -= 1 + fully_reduced_jacobian_indptr = np.unique(fully_reduced_jacobian_indptr) + + # Make sure two rows are removed and deduce shape of reduced jacobian + assert ( + len(fully_reduced_jacobian_indptr) == len(self.reduced_jacobian.indptr) - 2 + ), "Two rows should be removed." + fully_reduced_jacobian_shape = ( + len(fully_reduced_jacobian_indptr) - 1, + len(fully_reduced_jacobian_indptr) - 1, + ) - return flat_flux, flat_pressure, flat_lagrange_multiplier + # Cache the fully reduced jacobian + self.fully_reduced_jacobian = sps.csc_matrix( + ( + fully_reduced_jacobian_data, + fully_reduced_jacobian_indices, + fully_reduced_jacobian_indptr, + ), + shape=fully_reduced_jacobian_shape, + ) + """sps.csc_matrix: fully reduced jacobian""" + + # Cache the indices and indptr + self.fully_reduced_jacobian_indices = fully_reduced_jacobian_indices.copy() + """np.ndarray: indices of the fully reduced jacobian""" + + self.fully_reduced_jacobian_indptr = fully_reduced_jacobian_indptr.copy() + """np.ndarray: indptr of the fully reduced jacobian""" - # ! ---- Projections inbetween faces and cells ---- + self.fully_reduced_jacobian_shape = fully_reduced_jacobian_shape + """tuple: shape of the fully reduced jacobian""" + + # Step 4: Identify inclusions (index arrays) + + # Define reduced system indices wrt full system + reduced_system_indices = np.concatenate( + [self.potential_indices, self.lagrange_multiplier_indices] + ) + + # Define fully reduced system indices wrt reduced system - need to remove cell + # (and implicitly lagrange multiplier) + self.fully_reduced_system_indices = np.delete( + np.arange(self.grid.num_cells), self.constrained_cell_flat_index + ) + """np.ndarray: indices of the fully reduced system in terms of reduced system""" - def cell_reconstruction(self, flat_flux: np.ndarray) -> np.ndarray: - """Reconstruct the fluxes on the cells from the fluxes on the faces. + # Define fully reduced system indices wrt full system + self.fully_reduced_system_indices_full = reduced_system_indices[ + self.fully_reduced_system_indices + ] + """np.ndarray: indices of the fully reduced system in terms of full system""" + + def _setup_acceleration(self) -> None: + """Setup of acceleration methods.""" + + # ! ---- Acceleration ---- + aa_depth = self.options.get("aa_depth", 0) + aa_restart = self.options.get("aa_restart", None) + self.anderson = ( + darsia.AndersonAcceleration( + dimension=None, depth=aa_depth, restart=aa_restart + ) + if aa_depth > 0 + else None + ) + """darsia.AndersonAcceleration: Anderson acceleration""" + + # ! ---- Effective quantities ---- + + def compute_transport_density(self, solution: np.ndarray) -> np.ndarray: + """Compute the transport density from the solution. Args: - flat_flux (np.ndarray): flat fluxes (normal fluxes on the faces) + solution (np.ndarray): solution Returns: - np.ndarray: cell-based vectorial fluxes + np.ndarray: transport density """ - # TODO replace by sparse matrix multiplication + # Convert (scalar) normal fluxes to vector-valued fluxes on cells + flat_flux = solution[self.flux_slice] + cell_flux = darsia.face_to_cell(self.grid, flat_flux) + # Simply take the norm without any other integration + norm = np.linalg.norm(cell_flux, 2, axis=-1) + return norm + + def l1_dissipation(self, flat_flux: np.ndarray, mode: str) -> float: + """Compute the l1 dissipation potential of the solution. - # Reshape fluxes - use duality of faces and normals - horizontal_fluxes = flat_flux[: self.num_edges_axis[0]].reshape( - self.vertical_edges_shape - ) - vertical_fluxes = flat_flux[self.num_edges_axis[0] :].reshape( - self.horizontal_edges_shape - ) + Args: + flat_flux (np.ndarray): flat fluxes + + Returns: + float: l1 dissipation potential - # Determine a cell-based Raviart-Thomas reconstruction of the fluxes - cell_flux = np.zeros((*self.dim_cells, self.dim), dtype=float) - # Horizontal fluxes - cell_flux[:, :-1, 0] += 0.5 * horizontal_fluxes - cell_flux[:, 1:, 0] += 0.5 * horizontal_fluxes - # Vertical fluxes - cell_flux[:-1, :, 1] += 0.5 * vertical_fluxes - cell_flux[1:, :, 1] += 0.5 * vertical_fluxes + """ + if mode == "cell_arithmetic": + cell_flux = darsia.face_to_cell(self.grid, flat_flux) + cell_flux_norm = np.ravel(np.linalg.norm(cell_flux, 2, axis=-1)) + return self.mass_matrix_cells.dot(cell_flux_norm).sum() + elif mode == "face_arithmetic": + face_flux_norm = self.vector_face_flux_norm(flat_flux, "face_arithmetic") + return self.mass_matrix_faces.dot(face_flux_norm).sum() - return cell_flux + # ! ---- Lumping of effective mobility - def face_restriction(self, cell_flux: np.ndarray) -> np.ndarray: - """Restrict the fluxes on the cells to the faces. + def vector_face_flux_norm(self, flat_flux: np.ndarray, mode: str) -> np.ndarray: + """Compute the norm of the vector-valued fluxes on the faces. Args: - cell_flux (np.ndarray): cell-based fluxes + flat_flux (np.ndarray): flat fluxes (normal fluxes on the faces) + mode (str): mode of the norm, either "cell_arithmetic", "cell_harmonic" or + "face_arithmetic". In the cell-based modes, the fluxes are projected to + the cells and the norm is computed there. In the face-based mode, the + norm is computed directly on the faces. Returns: - np.ndarray: face-based fluxes + np.ndarray: norm of the vector-valued fluxes on the faces """ - # TODO replace by sparse matrix multiplication - # Determine the fluxes on the faces - horizontal_fluxes = 0.5 * (cell_flux[:, :-1, 0] + cell_flux[:, 1:, 0]) - vertical_fluxes = 0.5 * (cell_flux[:-1, :, 1] + cell_flux[1:, :, 1]) + # Determine the norm of the fluxes on the faces + if mode in ["cell_arithmetic", "cell_harmonic"]: + # Consider the piecewise constant projection of vector valued fluxes + cell_flux = darsia.face_to_cell(self.grid, flat_flux) + # Determine the norm of the fluxes on the cells + cell_flux_norm = np.maximum( + np.linalg.norm(cell_flux, 2, axis=-1), self.regularization + ) + # Determine averaging mode from mode - either arithmetic or harmonic + average_mode = mode.split("_")[1] + flat_flux_norm = darsia.cell_to_face( + self.grid, cell_flux_norm, mode=average_mode + ) - # Reshape the fluxes - flat_flux = np.concatenate( - [horizontal_fluxes.ravel(), vertical_fluxes.ravel()], axis=0 - ) + elif mode == "face_arithmetic": + # Define natural vector valued flux on faces (taking arithmetic averages + # of continuous fluxes over cells evaluated at faces) + full_face_flux = self.face_reconstruction(flat_flux) + # Determine the l2 norm of the fluxes on the faces + flat_flux_norm = np.linalg.norm(full_face_flux, 2, axis=1) + + else: + raise ValueError(f"Mode {mode} not supported.") - return flat_flux + return flat_flux_norm - def face_restriction_scalar(self, cell_qty: np.ndarray) -> np.ndarray: - """Restrict the fluxes on the cells to the faces. + # ! ---- Solver methods ---- + + def linear_solve( + self, + matrix: sps.csc_matrix, + rhs: np.ndarray, + previous_solution: Optional[np.ndarray] = None, + reuse_solver: bool = False, + ) -> tuple: + """Solve the linear system. + + For reusing the setup, the resulting solver is cached as self.linear_solver. Args: - cell_qty (np.ndarray): cell-based quantity + matrix (sps.csc_matrix): matrix + rhs (np.ndarray): right hand side + previous_solution (np.ndarray): previous solution. Defaults to None. Returns: - np.ndarray: face-based quantity + tuple: solution, stats """ - # Determine the fluxes on the faces - horizontal_face_qty = 0.5 * (cell_qty[:, :-1] + cell_qty[:, 1:]) - vertical_face_qty = 0.5 * (cell_qty[:-1, :] + cell_qty[1:, :]) + setup_linear_solver = not (reuse_solver) or not (hasattr(self, "linear_solver")) + + if self.linear_solver_type == "lu": + # Setup LU factorization for the full system + tic = time.time() + if setup_linear_solver: + self.linear_solver = sps.linalg.splu(matrix) + time_setup = time.time() - tic + + # Solve the full system + tic = time.time() + solution = self.linear_solver.solve(rhs) + time_solve = time.time() - tic + + elif self.linear_solver_type in [ + "lu-flux-reduced", + "amg-flux-reduced", + "lu-potential", + "amg-potential", + ]: + # Solve potential-multiplier problem + + # Allocate memory for solution + solution = np.zeros_like(rhs) + + # Reduce flux block + tic = time.time() + ( + self.reduced_matrix, + self.reduced_rhs, + matrix_flux_inv, + ) = self.remove_flux(matrix, rhs) + + if self.linear_solver_type == "lu-flux-reduced": + # LU factorization for reduced system + if setup_linear_solver: + self.linear_solver = sps.linalg.splu(self.reduced_matrix) + time_setup = time.time() - tic + + # Solve for the potential and lagrange multiplier + tic = time.time() + solution[self.reduced_system_slice] = self.linear_solver.solve( + self.reduced_rhs + ) - # Reshape the fluxes - hardcoding the connectivity here - face_qty = np.concatenate( - [horizontal_face_qty.ravel(), vertical_face_qty.ravel()] - ) + elif self.linear_solver_type == "amg-flux-reduced": + # AMG solver for reduced system + if setup_linear_solver: + self.linear_solver = pyamg.smoothed_aggregation_solver( + self.reduced_matrix, **self.ml_options + ) + time_setup = time.time() - tic + + # Solve for the potential and lagrange multiplier + tic = time.time() + solution[self.reduced_system_slice] = self.linear_solver.solve( + self.reduced_rhs, + tol=self.tol_amg, + residuals=self.res_history_amg, + ) + else: + # Solve pure potential problem + + # NOTE: It is implicitly assumed that the lagrange multiplier is zero + # in the constrained cell. This is not checked here. And no update is + # performed. + if previous_solution is not None and ( + abs( + previous_solution[ + self.grid.num_faces + self.constrained_cell_flat_index + ] + ) + > 1e-6 + ): + raise NotImplementedError( + "Implementation requires solution satisfy the constraint." + ) - return face_qty + # Reduce to pure potential system + ( + self.fully_reduced_matix, + self.fully_reduced_rhs, + ) = self.remove_lagrange_multiplier( + self.reduced_matrix, + self.reduced_rhs, + ) - # ! ---- Effective quantities ---- + if self.linear_solver_type == "lu-potential": + # Finish LU factorization of the pure potential system + if setup_linear_solver: + self.linear_solver = sps.linalg.splu(self.fully_reduced_matrix) + time_setup = time.time() - tic + + # Solve the pure potential system + tic = time.time() + solution[ + self.fully_reduced_system_indices_full + ] = self.linear_solver.solve(self.fully_reduced_rhs) + + elif self.linear_solver_type == "amg-potential": + # Finish AMG setup of th pure potential system + if setup_linear_solver: + self.linear_solver = pyamg.smoothed_aggregation_solver( + self.fully_reduced_jacobian, **self.ml_options + ) + time_setup = time.time() - tic + + # Solve the pure potential system + tic = time.time() + solution[ + self.fully_reduced_system_indices_full + ] = self.linear_solver.solve( + self.fully_reduced_rhs, + tol=self.tol_amg, + residuals=self.res_history_amg, + ) - def effective_mobility(self, flat_flux: np.ndarray) -> np.ndarray: - """Compute the effective mobility of the solution. + # Compute flux update + solution[self.flux_slice] = matrix_flux_inv.dot( + rhs[self.flux_slice] + self.DT.dot(solution[self.reduced_system_slice]) + ) + time_solve = time.time() - tic + + stats = { + "time setup": time_setup, + "time solve": time_solve, + } + if self.linear_solver_type in ["amg-flux-reduced", "amg-potential"]: + stats["amg residuals"] = self.res_history_amg + stats["amg num iterations"] = len(self.res_history_amg) + stats["amg residual"] = self.res_history_amg[-1] + + return solution, stats + + def remove_flux(self, jacobian: sps.csc_matrix, residual: np.ndarray) -> tuple: + """Remove the flux block from the jacobian and residual. Args: - flat_flux (np.ndarray): flat fluxes + jacobian (sps.csc_matrix): jacobian + residual (np.ndarray): residual Returns: - np.ndarray: effective mobility + tuple: reduced jacobian, reduced residual, inverse of flux block + """ - # TODO Use improved quadrature? - cell_flux = self.cell_reconstruction(flat_flux) - return np.linalg.norm(cell_flux, 2, axis=-1) + # Build Schur complement wrt flux-block + J_inv = sps.diags(1.0 / jacobian.diagonal()[self.flux_slice]) + schur_complement = self.D.dot(J_inv.dot(self.DT)) - def l1_dissipation(self, solution: np.ndarray) -> float: - """Compute the l1 dissipation potential of the solution. + # Gauss eliminiation on matrices + reduced_jacobian = self.jacobian_subblock + schur_complement + + # Gauss elimination on vectors + reduced_residual = residual[self.reduced_system_slice].copy() + reduced_residual -= self.D.dot(J_inv.dot(residual[self.flux_slice])) + + return reduced_jacobian, reduced_residual, J_inv + + def remove_lagrange_multiplier(self, reduced_jacobian, reduced_residual) -> tuple: + """Shortcut for removing the lagrange multiplier from the reduced jacobian. Args: - solution (np.ndarray): solution + reduced_jacobian (sps.csc_matrix): reduced jacobian + reduced_residual (np.ndarray): reduced residual Returns: - float: l1 dissipation potential + tuple: fully reduced jacobian, fully reduced residual """ - # TODO use improved quadrature? - flat_flux, _, _ = self.split_solution(solution) - cell_flux = self.cell_reconstruction(flat_flux) - return np.sum(np.prod(self.voxel_size) * np.linalg.norm(cell_flux, 2, axis=-1)) + # Make sure the jacobian is a CSC matrix + assert isinstance( + reduced_jacobian, sps.csc_matrix + ), "Jacobian should be a CSC matrix." + + # Effective Gauss-elimination for the particular case of the lagrange multiplier + self.fully_reduced_jacobian.data[:] = np.delete( + reduced_jacobian.data.copy(), self.rm_indices + ) + # NOTE: The indices have to be restored if the LU factorization is to be used + # FIXME omit if not required + self.fully_reduced_jacobian.indices = self.fully_reduced_jacobian_indices.copy() + + # Rhs is not affected by Gauss elimination as it is assumed that the residual + # is zero in the constrained cell, and the pressure is zero there as well. + # If not, we need to do a proper Gauss elimination on the right hand side! + if abs(reduced_residual[-1]) > 1e-6: + raise NotImplementedError("Implementation requires residual to be zero.") + fully_reduced_residual = reduced_residual[ + self.fully_reduced_system_indices + ].copy() + + return self.fully_reduced_jacobian, fully_reduced_residual # ! ---- Main methods ---- @@ -448,88 +729,174 @@ def __call__( # Determine difference of distriutions and define corresponding rhs mass_diff = img_1.img - img_2.img flat_mass_diff = np.ravel(mass_diff) - self._problem_specific_setup(mass_diff) # Main method distance, solution, status = self._solve(flat_mass_diff) # Split the solution - flat_flux, flat_pressure, _ = self.split_solution(solution) + flat_flux = solution[self.flux_slice] + flat_potential = solution[self.potential_slice] - # Reshape the fluxes and pressure - flux = self.cell_reconstruction(flat_flux) - pressure = flat_pressure.reshape(self.dim_cells) + # Reshape the fluxes and potential to grid format + flux = darsia.face_to_cell(self.grid, flat_flux) + potential = flat_potential.reshape(self.grid.shape) - # Determine effective mobility - mobility = self.effective_mobility(flat_flux) + # Determine transport density + transport_density = self.compute_transport_density(solution) # Stop taking time toc = time.time() status["elapsed_time"] = toc - tic + print("Elapsed time: ", toc - tic) # Plot the solution if plot_solution: - self._plot_solution(mass_diff, flux, pressure, mobility) + self._plot_solution(mass_diff, flux, potential, transport_density) if return_solution: - return distance, flux, pressure, mobility, status + return distance, flux, potential, transport_density, status else: return distance + # TODO rm. def _plot_solution( self, mass_diff: np.ndarray, flux: np.ndarray, - pressure: np.ndarray, - mobility: np.ndarray, + potential: np.ndarray, + transport_density: np.ndarray, ) -> None: + """Plot the solution. + + Args: + mass_diff (np.ndarray): difference of mass distributions + flux (np.ndarray): fluxes + potential (np.ndarray): potential + transport_density (np.ndarray): transport density + + Raises: + NotImplementedError: plotting only implemented for 2D + + """ + + if self.grid.dim != 2: + raise NotImplementedError("Plotting only implemented for 2D.") + + # Fetch options + plot_options = self.options.get("plot_options", {}) + name = plot_options.get("name", None) + + # Store plot + save_plot = plot_options.get("save", False) + if save_plot: + folder = plot_options.get("folder", ".") + Path(folder).mkdir(parents=True, exist_ok=True) + show_plot = plot_options.get("show", True) + + # Control of flux arrows + scaling = plot_options.get("scaling", 1.0) + resolution = plot_options.get("resolution", 1) + # Meshgrid Y, X = np.meshgrid( - self.voxel_size[0] * (0.5 + np.arange(self.shape[0] - 1, -1, -1)), - self.voxel_size[1] * (0.5 + np.arange(self.shape[1])), + self.voxel_size[0] * (0.5 + np.arange(self.grid.shape[0] - 1, -1, -1)), + self.voxel_size[1] * (0.5 + np.arange(self.grid.shape[1])), indexing="ij", ) - scaling = self.options.get("scaling", 1.0) - - # Plot the fluxes and pressure - plt.figure("Beckman solution") - plt.pcolormesh(X, Y, pressure, cmap="turbo") - plt.colorbar() + # Plot the potential + plt.figure("Beckman solution potential") + plt.pcolormesh(X, Y, potential, cmap="turbo") + plt.colorbar(label="potential") plt.quiver( - X, - Y, - scaling * flux[:, :, 0], - -scaling * flux[:, :, 1], + X[::resolution, ::resolution], + Y[::resolution, ::resolution], + scaling * flux[::resolution, ::resolution, 0], + -scaling * flux[::resolution, ::resolution, 1], angles="xy", scale_units="xy", scale=1, - alpha=0.5, + alpha=0.25, + width=0.005, ) + plt.xlabel("x [m]") + plt.ylabel("y [m]") + # plt.ylim(top=0.08) # TODO rm? + if save_plot: + plt.savefig( + folder + "/" + name + "_beckman_solution_potential.png", + dpi=500, + transparent=True, + ) + # Plot the fluxes plt.figure("Beckman solution fluxes") - plt.pcolormesh(X, Y, mass_diff, cmap="turbo") - plt.colorbar() + plt.pcolormesh(X, Y, mass_diff, cmap="turbo") # , vmin=-1, vmax=3.5) + plt.colorbar(label="mass difference") plt.quiver( - X, - Y, - scaling * flux[:, :, 0], - -scaling * flux[:, :, 1], + X[::resolution, ::resolution], + Y[::resolution, ::resolution], + scaling * flux[::resolution, ::resolution, 0], + -scaling * flux[::resolution, ::resolution, 1], angles="xy", scale_units="xy", scale=1, - alpha=0.5, + alpha=0.25, + width=0.005, ) + plt.xlabel("x [m]") + plt.ylabel("y [m]") + # plt.ylim(top=0.08) + plt.text( + 0.0025, + 0.075, + name, + color="white", + alpha=0.9, + rotation=0, + fontsize=14, + ) # TODO rm? + if save_plot: + plt.savefig( + folder + "/" + name + "_beckman_solution_fluxes.png", + dpi=500, + transparent=True, + ) - plt.figure("Beckman solution mobility") - plt.pcolormesh(X, Y, mobility, cmap="turbo") - plt.colorbar() + # Plot the transport density + plt.figure("L1 optimal transport density") + plt.pcolormesh(X, Y, transport_density, cmap="turbo") + plt.colorbar(label="flux modulus") + plt.xlabel("x [m]") + plt.ylabel("y [m]") + # plt.ylim(top=0.08) # TODO rm? + if save_plot: + plt.savefig( + folder + "/" + name + "_beckman_solution_transport_density.png", + dpi=500, + transparent=True, + ) - plt.show() + if show_plot: + plt.show() + else: + plt.close("all") class WassersteinDistanceNewton(VariationalWassersteinDistance): - """Class to determine the L1 EMD/Wasserstein distance solved with Newton's method.""" + """Class to determine the L1 EMD/Wasserstein distance solved with Newton's method. + + Here, self.L has the interpretation of a lower cut-off value in the linearization + only. With such relaxation, the BEckman problem itself is not regularized, but + instead the solution trajectory is merely affected. + + """ + + def __init__(self, grid, options) -> None: + super().__init__(grid, options) + + self.L = self.options.get("L", 1.0) + """float: relaxation parameter, lower cut-off for the mobility""" def residual(self, rhs: np.ndarray, solution: np.ndarray) -> np.ndarray: """Compute the residual of the solution. @@ -542,20 +909,20 @@ def residual(self, rhs: np.ndarray, solution: np.ndarray) -> np.ndarray: np.ndarray: residual """ - flat_flux, _, _ = self.split_solution(solution) - cell_flux = self.cell_reconstruction(flat_flux) - cell_flux_norm = np.maximum( - np.linalg.norm(cell_flux, 2, axis=-1), self.regularization + flat_flux = solution[self.flux_slice] + mode = self.options.get("mode", "face_arithmetic") + flat_flux_norm = np.maximum( + self.vector_face_flux_norm(flat_flux, mode=mode), self.regularization ) - cell_flux_normed = cell_flux / cell_flux_norm[..., None] - flat_flux_normed = self.face_restriction(cell_flux_normed) + flat_flux_normed = flat_flux / flat_flux_norm + return ( rhs - self.broken_darcy.dot(solution) - - self.flux_embedding.dot(self.mass_matrix_edges.dot(flat_flux_normed)) + - self.flux_embedding.dot(self.mass_matrix_faces.dot(flat_flux_normed)) ) - def jacobian_lu(self, solution: np.ndarray) -> sps.linalg.splu: + def jacobian(self, solution: np.ndarray) -> sps.linalg.LinearOperator: """Compute the LU factorization of the jacobian of the solution. Args: @@ -565,43 +932,53 @@ def jacobian_lu(self, solution: np.ndarray) -> sps.linalg.splu: sps.linalg.splu: LU factorization of the jacobian """ - flat_flux, _, _ = self.split_solution(solution) - cell_flux = self.cell_reconstruction(flat_flux) - self.regularization = self.options.get("regularization", 0.0) - cell_flux_norm = np.maximum( - np.linalg.norm(cell_flux, 2, axis=-1), self.regularization + flat_flux = solution[self.flux_slice] + mode = self.options.get("mode", "face_arithmetic") + flat_flux_norm = np.maximum( + self.vector_face_flux_norm(flat_flux, mode=mode), self.regularization ) - flat_flux_norm = self.face_restriction_scalar(cell_flux_norm) approx_jacobian = sps.bmat( [ [ sps.diags(np.maximum(self.L, 1.0 / flat_flux_norm), dtype=float) - * self.mass_matrix_edges, + * self.mass_matrix_faces, -self.div.T, None, ], - [self.div, None, -self.pressure_constraint.T], - [None, self.pressure_constraint, None], + [self.div, None, -self.potential_constraint.T], + [None, self.potential_constraint, None], ], format="csc", ) - approx_jacobian_lu = sps.linalg.splu(approx_jacobian) - return approx_jacobian_lu + return approx_jacobian - def _solve(self, flat_mass_diff): - # Observation: AA can lead to less stagnation, more accurate results, and therefore - # better solutions to mu and u. Higher depth is better, but more expensive. + def _solve(self, flat_mass_diff: np.ndarray) -> tuple[float, np.ndarray, dict]: + """Solve the Beckman problem using Newton's method. + + Args: + flat_mass_diff (np.ndarray): difference of mass distributions + + Returns: + tuple: distance, solution, status + + """ + # TODO rm: Observation: AA can lead to less stagnation, more accurate results, + # and therefore better solutions to mu and u. Higher depth is better, but more + # expensive. + + # Setup + tic = time.time() # Solver parameters num_iter = self.options.get("num_iter", 100) - tol = self.options.get("tol", 1e-6) + tol_residual = self.options.get("tol_residual", 1e-6) + tol_increment = self.options.get("tol_increment", 1e-6) tol_distance = self.options.get("tol_distance", 1e-6) - self.L = self.options.get("L", 1.0) # Define right hand side rhs = np.concatenate( [ - np.zeros(self.num_edges, dtype=float), + np.zeros(self.grid.num_faces, dtype=float), self.mass_matrix_cells.dot(flat_mass_diff), np.zeros(1, dtype=float), ] @@ -610,64 +987,145 @@ def _solve(self, flat_mass_diff): # Initialize solution solution_i = np.zeros_like(rhs) + # Initialize container for storing the convergence history + convergence_history = { + "distance": [], + "residual": [], + "decomposed residual": [], + "increment": [], + "decomposed increment": [], + "distance increment": [], + "timing": [], + } + + # Print header + if self.verbose: + print( + "--- ; ", + "Newton iteration", + "distance", + "residual", + "mass conservation residual", + "increment", + "flux increment", + "distance increment", + ) + # Newton iteration for iter in range(num_iter): # Keep track of old flux, and old distance old_solution_i = solution_i.copy() - old_distance = self.l1_dissipation(solution_i) + old_flux = solution_i[self.flux_slice] + old_distance = self.l1_dissipation(old_flux, "cell_arithmetic") - # Newton step + # Assemble linear problem in Newton step + tic = time.time() if iter == 0: - residual_i = ( - rhs.copy() - ) # Aim at Darcy-like initial guess after first iteration. + # Determine residual and (full) Jacobian of a linear Darcy problem + residual_i = rhs.copy() + approx_jacobian = self.darcy_init.copy() else: + # Determine residual and (full) Jacobian residual_i = self.residual(rhs, solution_i) - jacobian_lu = self.jacobian_lu(solution_i) - update_i = jacobian_lu.solve(residual_i) + approx_jacobian = self.jacobian(solution_i) + toc = time.time() + time_assemble = toc - tic + + # Solve linear system for the update + update_i, stats_i = self.linear_solve( + approx_jacobian, residual_i, solution_i + ) + + # Diagnostics + # TODO move? + if self.linear_solver_type in ["amg-flux-reduced", "amg-potential"]: + if self.options.get("linear_solver_verbosity", False): + # print(ml) # TODO rm? + print( + f"""AMG iterations: {stats_i["amg num iterations"]}; """ + f"""Residual after AMG step: {stats_i["amg residual"]}""" + ) + + # Update the solution with the full Netwon step solution_i += update_i # Apply Anderson acceleration to flux contribution (the only nonlinear part). + # Application to full solution, or just the potential, lead to divergence, + # while application to the flux, results in improved performance. + tic = time.time() if self.anderson is not None: - solution_i[: self.num_edges] = self.anderson( - solution_i[: self.num_edges], update_i[: self.num_edges], iter + solution_i[self.flux_slice] = self.anderson( + solution_i[self.flux_slice], + update_i[self.flux_slice], + iter, ) - # TODO try for full solution + toc = time.time() + time_anderson = toc - tic + + # Update stats + stats_i["time assemble"] = time_assemble + stats_i["time acceleration"] = time_anderson # Update distance - new_distance = self.l1_dissipation(solution_i) + new_flux = solution_i[self.flux_slice] + new_distance = self.l1_dissipation(new_flux, "cell_arithmetic") # Compute the error: - # - residual + # - full residual + # - residual of the flux equation # - residual of mass conservation equation - # - increment + # - residual of the constraint equation + # - full increment # - flux increment + # - potential increment + # - lagrange multiplier increment + # - distance increment + increment = solution_i - old_solution_i error = [ np.linalg.norm(residual_i, 2), - np.linalg.norm(residual_i[self.num_edges : -1], 2), - np.linalg.norm(solution_i - old_solution_i, 2), - np.linalg.norm((solution_i - old_solution_i)[: self.num_edges], 2), + [ + np.linalg.norm(residual_i[self.flux_slice], 2), + np.linalg.norm(residual_i[self.potential_slice], 2), + np.linalg.norm(residual_i[self.lagrange_multiplier_slice], 2), + ], + np.linalg.norm(increment, 2), + [ + np.linalg.norm(increment[self.flux_slice], 2), + np.linalg.norm(increment[self.potential_slice], 2), + np.linalg.norm(increment[self.lagrange_multiplier_slice], 2), + ], + abs(new_distance - old_distance), ] + # Update convergence history + convergence_history["distance"].append(new_distance) + convergence_history["residual"].append(error[0]) + convergence_history["decomposed residual"].append(error[1]) + convergence_history["increment"].append(error[2]) + convergence_history["decomposed increment"].append(error[3]) + convergence_history["distance increment"].append(error[4]) + convergence_history["timing"].append(stats_i) + + new_distance_faces = self.l1_dissipation(new_flux, "face_arithmetic") if self.verbose: print( "Newton iteration", iter, new_distance, - old_distance - new_distance, + new_distance_faces, error[0], # residual error[1], # mass conservation residual error[2], # full increment error[3], # flux increment + error[4], # distance increment + stats_i, # timing ) # Stopping criterion # TODO include criterion build on staganation of the solution - # TODO include criterion on distance. - if ( - iter > 1 - and min([error[0], error[2]]) < tol - or abs(new_distance - old_distance) < tol_distance + if iter > 1 and ( + (error[0] < tol_residual and error[2] < tol_increment) + or error[4] < tol_distance # TODO rm the latter ): break @@ -681,34 +1139,93 @@ def _solve(self, flat_mass_diff): "increment": error[2], "flux increment": error[3], "distance increment": abs(new_distance - old_distance), + "convergence history": convergence_history, } return new_distance, solution_i, status class WassersteinDistanceBregman(VariationalWassersteinDistance): - def _problem_specific_setup(self, mass_diff: np.ndarray) -> None: - super()._problem_specific_setup(mass_diff) + # TODO __init__ correct method? + def __init__( + self, + grid: darsia.Grid, + options: dict = {}, + ) -> None: + super().__init__(grid, options) self.L = self.options.get("L", 1.0) - l_scheme_mixed_darcy = sps.bmat( - [ - [self.L * self.mass_matrix_edges, -self.div.T, None], - [self.div, None, -self.pressure_constraint.T], - [None, self.pressure_constraint, None], - ], - format="csc", - ) - self.l_scheme_mixed_darcy_lu = sps.linalg.splu(l_scheme_mixed_darcy) + """Penality parameter for the Bregman iteration.""" + + self.force_slice = slice(self.grid.num_faces, None) + """slice: slice for the force.""" + + def _shrink( + self, + flat_flux: np.ndarray, + shrink_factor: Union[float, np.ndarray], + mode: str = "cell_arithmetic", + ) -> np.ndarray: + """Shrink operation in the split Bregman method. + + Operation on fluxes. + + Args: + flat_flux (np.ndarray): flux + shrink_factor (float or np.ndarray): shrink factor + mode (str, optional): mode of the shrink operation. Defaults to "cell_arithmetic". + + Returns: + np.ndarray: shrunk fluxes + + """ + if mode == "cell_arithmetic": + # Idea: Determine the shrink factor based on the cell reconstructions of the + # fluxes. Convert cell-based shrink factors to face-based shrink factors + # through arithmetic averaging. + cell_flux = darsia.face_to_cell(self.grid, flat_flux) + norm = np.linalg.norm(cell_flux, 2, axis=-1) + cell_scaling = np.maximum(norm - shrink_factor, 0) / ( + norm + self.regularization + ) + flat_scaling = darsia.cell_to_face( + self.grid, cell_scaling, mode="arithmetic" + ) + + elif mode == "face_arithmetic": + # Define natural vector valued flux on faces (taking arithmetic averages + # of continuous fluxes over cells evaluated at faces) + full_face_flux = self.face_reconstruction(flat_flux) + # Determine the l2 norm of the fluxes on the faces, add some regularization + norm = np.linalg.norm(full_face_flux, 2, axis=1) + flat_scaling = np.maximum(norm - shrink_factor, 0) / ( + norm + self.regularization + ) + + elif mode == "face_normal": + # Only consider normal direction (does not take into account the full flux) + # TODO rm. + norm = np.linalg.norm(flat_flux, 2, axis=-1) + flat_scaling = np.maximum(norm - shrink_factor, 0) / ( + norm + self.regularization + ) + + else: + raise NotImplementedError(f"Mode {mode} not supported.") + + return flat_scaling * flat_flux def _solve(self, flat_mass_diff): # Solver parameters num_iter = self.options.get("num_iter", 100) - # tol = self.options.get("tol", 1e-6) # TODO make use of tol, or remove - self.L = self.options.get("L", 1.0) + tol_residual = self.options.get("tol_residual", 1e-6) + tol_increment = self.options.get("tol_increment", 1e-6) + tol_distance = self.options.get("tol_distance", 1e-6) + # Relaxation parameter + self.L = self.options.get("L", 1.0) rhs = np.concatenate( [ - np.zeros(self.num_edges, dtype=float), + np.zeros(self.grid.num_faces, dtype=float), self.mass_matrix_cells.dot(flat_mass_diff), np.zeros(1, dtype=float), ] @@ -717,58 +1234,263 @@ def _solve(self, flat_mass_diff): # Keep track of how often the distance increases. num_neg_diff = 0 - # Bregman iterations - solution_i = np.zeros_like(rhs) - for iter in range(num_iter): - old_distance = self.l1_dissipation(solution_i) - - # 1. Solve linear system with trust in current flux. - flat_flux_i, _, _ = self.split_solution(solution_i) - rhs_i = rhs.copy() - rhs_i[: self.num_edges] = self.L * self.mass_matrix_edges.dot(flat_flux_i) - intermediate_solution_i = self.l_scheme_mixed_darcy_lu.solve(rhs_i) - - # 2. Shrink step for vectorial fluxes. To comply with the RT0 setting, the - # shrinkage operation merely determines the scalar. We still aim at - # following along the direction provided by the vectorial fluxes. - intermediate_flat_flux_i, _, _ = self.split_solution( - intermediate_solution_i - ) - # new_flat_flux_i = np.sign(intermediate_flat_flux_i) * ( - # np.maximum(np.abs(intermediate_flat_flux_i) - 1.0 / self.L, 0.0) - # ) - cell_intermediate_flux_i = self.cell_reconstruction( - intermediate_flat_flux_i - ) - norm = np.linalg.norm(cell_intermediate_flux_i, 2, axis=-1) - cell_scaling = np.maximum(norm - 1 / self.L, 0) / ( - norm + self.regularization + # Initialize container for storing the convergence history + convergence_history = { + "distance": [], + "mass residual": [], + "force": [], + "flux increment": [], + "aux increment": [], + "force increment": [], + "distance increment": [], + "timing": [], + } + + # Print header + if self.verbose: + print( + "--- ; ", + "Bregman iteration", + "L", + "distance", + "mass conservation residual", + "[flux, aux, force] increment", + "distance increment", ) - flat_scaling = self.face_restriction_scalar(cell_scaling) - new_flat_flux_i = flat_scaling * intermediate_flat_flux_i - # Apply Anderson acceleration to flux contribution (the only nonlinear part). - if self.anderson is not None: - flux_inc = new_flat_flux_i - flat_flux_i - new_flat_flux_i = self.anderson(new_flat_flux_i, flux_inc, iter) + # Initialize Bregman variables and flux with Darcy flow + shrink_mode = "face_arithmetic" + dissipation_mode = "cell_arithmetic" + weight = self.L + shrink_factor = 1.0 / self.L + + # Solve linear Darcy problem as initial guess + l_scheme_mixed_darcy = sps.bmat( + [ + [self.L * self.mass_matrix_faces, -self.div.T, None], + [self.div, None, -self.potential_constraint.T], + [None, self.potential_constraint, None], + ], + format="csc", + ) + solution_i = np.zeros_like(rhs, dtype=float) + solution_i, _ = self.linear_solve(l_scheme_mixed_darcy, rhs, solution_i) + + # Extract intial values + old_flux = solution_i[self.flux_slice] + old_aux_flux = self._shrink(old_flux, shrink_factor, shrink_mode) + old_force = old_flux - old_aux_flux + old_distance = self.l1_dissipation(old_flux, dissipation_mode) + + for iter in range(num_iter): + bregman_mode = self.options.get("bregman_mode", "standard") + if bregman_mode == "standard": + # std split Bregman method + + # 1. Make relaxation step (solve quadratic optimization problem) + tic = time.time() + rhs_i = rhs.copy() + rhs_i[self.flux_slice] = self.L * self.mass_matrix_faces.dot( + old_aux_flux - old_force + ) + solution_i, _ = self.linear_solve( + l_scheme_mixed_darcy, rhs_i, reuse_solver=True + ) + new_flux = solution_i[self.flux_slice] + time_linearization = time.time() - tic + + # 2. Shrink step for vectorial fluxes. To comply with the RT0 setting, the + # shrinkage operation merely determines the scalar. We still aim at + # following along the direction provided by the vectorial fluxes. + tic = time.time() + new_aux_flux = self._shrink( + new_flux + old_force, shrink_factor, shrink_mode + ) + time_shrink = time.time() - tic + + # 3. Update force + new_force = old_force + new_flux - new_aux_flux + + # Apply Anderson acceleration to flux contribution (the only nonlinear part). + tic = time.time() + if self.anderson is not None: + aux_inc = new_aux_flux - old_aux_flux + force_inc = new_force - old_force + inc = np.concatenate([aux_inc, force_inc]) + iteration = np.concatenate([new_aux_flux, new_force]) + new_iteration = self.anderson(iteration, inc, iter) + new_aux_flux = new_iteration[self.flux_slice] + new_force = new_iteration[self.force_slice] + + toc = time.time() + time_anderson = toc - tic + + elif bregman_mode == "reordered": + # Reordered split Bregman method + + # 1. Shrink step for vectorial fluxes. To comply with the RT0 setting, the + # shrinkage operation merely determines the scalar. We still aim at + # following along the direction provided by the vectorial fluxes. + tic = time.time() + new_aux_flux = self._shrink( + old_flux + old_force, shrink_factor, shrink_mode + ) + time_shrink = time.time() - tic + + # 2. Update force + new_force = old_force + old_flux - new_aux_flux + + # 3. Solve linear system with trust in current flux. + tic = time.time() + rhs_i = rhs.copy() + rhs_i[self.flux_slice] = self.L * self.mass_matrix_faces.dot( + new_aux_flux - new_force + ) + solution_i, _ = self.linear_solve( + l_scheme_mixed_darcy, rhs_i, reuse_solver=True + ) + new_flux = solution_i[self.flux_slice] + time_linearization = time.time() - tic + + # Apply Anderson acceleration to flux contribution (the only nonlinear part). + tic = time.time() + if self.anderson is not None: + flux_inc = new_flux - old_flux + force_inc = new_force - old_force + inc = np.concatenate([flux_inc, force_inc]) + iteration = np.concatenate([new_flux, new_force]) + # new_flux = self.anderson(new_flux, flux_inc, iter) + new_iteration = self.anderson(iteration, inc, iter) + new_flux = new_iteration[self.flux_slice] + new_force = new_iteration[self.force_slice] + toc = time.time() + time_anderson = toc - tic + + elif bregman_mode == "adaptive": + # Bregman split with updated weight + update_cond = self.options.get( + "bregman_update_cond", lambda iter: False + ) + update_solver = update_cond(iter) + if update_solver: + # TODO: self._update_weight(old_flux) + + # Update weight as the inverse of the norm of the flux + old_flux_norm = np.maximum( + self.vector_face_flux_norm(old_flux, "face_arithmetic"), + self.regularization, + ) + old_flux_norm_inv = 1.0 / old_flux_norm + weight = sps.diags(old_flux_norm_inv) + shrink_factor = old_flux_norm + + # Redefine Darcy system + l_scheme_mixed_darcy = sps.bmat( + [ + [weight * self.mass_matrix_faces, -self.div.T, None], + [self.div, None, -self.potential_constraint.T], + [None, self.potential_constraint, None], + ], + format="csc", + ) + + # 1. Make relaxation step (solve quadratic optimization problem) + tic = time.time() + rhs_i = rhs.copy() + rhs_i[self.flux_slice] = weight * self.mass_matrix_faces.dot( + old_aux_flux - old_force + ) + solution_i, _ = self.linear_solve( + l_scheme_mixed_darcy, rhs_i, reuse_solver=not (update_solver) + ) + + # Diagnostics + if self.linear_solver_type in ["amg-flux-reduced", "amg-potential"]: + if self.options.get("linear_solver_verbosity", False): + num_amg_iter = len(self.res_history_amg) + res_amg = self.res_history_amg[-1] + print(self.l_scheme_mixed_darcy_solver) + print( + f"""#AMG iterations: {num_amg_iter}; Residual after """ + f"""AMG step: {res_amg}""" + ) + + new_flux = solution_i[self.flux_slice] + time_linearization = time.time() - tic + + # 2. Shrink step for vectorial fluxes. To comply with the RT0 setting, the + # shrinkage operation merely determines the scalar. We still aim at + # following along the direction provided by the vectorial fluxes. + tic = time.time() + new_aux_flux = self._shrink( + new_flux + old_force, shrink_factor, shrink_mode + ) + time_shrink = time.time() - tic - # Measure error in terms of the increment of the flux - flux_diff = np.linalg.norm(new_flat_flux_i - flat_flux_i, 2) + # 3. Update force + new_force = old_force + new_flux - new_aux_flux - # Update flux solution - solution_i = intermediate_solution_i.copy() - solution_i[: self.num_edges] = new_flat_flux_i + # Apply Anderson acceleration to flux contribution (the only nonlinear part). + tic = time.time() + if self.anderson is not None: + aux_inc = new_aux_flux - old_aux_flux + force_inc = new_force - old_force + inc = np.concatenate([aux_inc, force_inc]) + iteration = np.concatenate([new_aux_flux, new_force]) + new_iteration = self.anderson(iteration, inc, iter) + new_aux_flux = new_iteration[self.flux_slice] + new_force = new_iteration[self.force_slice] + + toc = time.time() + time_anderson = toc - tic + + else: + raise NotImplementedError(f"Bregman mode {bregman_mode} not supported.") + + # Collect stats + stats_i = [time_linearization, time_shrink, time_anderson] # Update distance - new_distance = self.l1_dissipation(solution_i) + new_distance = self.l1_dissipation(new_flux, dissipation_mode) # Determine the error in the mass conservation equation mass_conservation_residual = np.linalg.norm( - (rhs_i - self.broken_darcy.dot(solution_i))[self.num_edges : -1], 2 + self.div.dot(new_flux) - rhs[self.potential_slice], 2 ) - # TODO include criterion build on staganation of the solution - # TODO include criterion on distance. + # Determine increments + flux_increment = new_flux - old_flux + aux_increment = new_aux_flux - old_aux_flux + force_increment = new_force - old_force + + # Determine force + force = np.linalg.norm(new_force, 2) + + # Compute the error: + # - residual of mass conservation equation - zero only if exact solver used + # - force + # - flux increment + # - aux increment + # - force increment + # - distance increment + error = [ + mass_conservation_residual, + force, + np.linalg.norm(flux_increment), + np.linalg.norm(aux_increment), + np.linalg.norm(force_increment), + abs(new_distance - old_distance), + ] + + # Update convergence history + convergence_history["distance"].append(new_distance) + convergence_history["mass residual"].append(error[0]) + convergence_history["force"].append(error[1]) + convergence_history["flux increment"].append(error[2]) + convergence_history["aux increment"].append(error[3]) + convergence_history["force increment"].append(error[4]) + convergence_history["distance increment"].append(error[5]) + convergence_history["timing"].append(stats_i) # Print status if self.verbose: @@ -776,60 +1498,88 @@ def _solve(self, flat_mass_diff): "Bregman iteration", iter, new_distance, - old_distance - new_distance, self.L, - flux_diff, - mass_conservation_residual, + error[0], # mass conservation residual + [ + error[2], # flux increment + error[3], # aux increment + error[4], # force increment + ], + error[5], # distance increment + # stats_i, # timings ) - ## Check stopping criterion # TODO. What is a good stopping criterion? - # if iter > 1 and (flux_diff < tol or mass_conservation_residual < tol: - # break - # Keep track if the distance increases. if new_distance > old_distance: num_neg_diff += 1 - # Increase L if stagnating of the distance increases too often. - # TODO restart anderson acceleration - update_l = self.options.get("update_l", True) - if update_l: - tol_distance = self.options.get("tol_distance", 1e-12) - max_iter_increase_diff = self.options.get("max_iter_increase_diff", 20) - l_factor = self.options.get("l_factor", 2) - if ( - abs(new_distance - old_distance) < tol_distance - or num_neg_diff > max_iter_increase_diff - ): - # Update L - self.L = self.L * l_factor - - # Update linear system - l_scheme_mixed_darcy = sps.bmat( - [ - [self.L * self.mass_matrix_edges, -self.div.T, None], - [self.div, None, -self.pressure_constraint.T], - [None, self.pressure_constraint, None], - ], - format="csc", - ) - self.l_scheme_mixed_darcy_lu = sps.linalg.splu(l_scheme_mixed_darcy) - - # Reset stagnation counter - num_neg_diff = 0 + # TODO include criterion build on staganation of the solution + if iter > 1 and ( + (error[0] < tol_residual and error[4] < tol_increment) + or error[5] < tol_distance + ): + break - L_max = self.options.get("L_max", 1e8) - if self.L > L_max: - break + # TODO rm? + # update_l = self.options.get("update_l", False) + # if update_l: + # tol_distance = self.options.get("tol_distance", 1e-12) + # max_iter_increase_diff = self.options.get( + # "max_iter_increase_diff", + # 20 + # ) + # l_factor = self.options.get("l_factor", 2) + # if ( + # abs(new_distance - old_distance) < tol_distance + # or num_neg_diff > max_iter_increase_diff + # ): + # # Update L + # self.L = self.L * l_factor + # + # # Update linear system + # l_scheme_mixed_darcy = sps.bmat( + # [ + # [ + # self.L * self.mass_matrix_faces, + # -self.div.T, + # None + # ], + # [self.div, None, -self.potential_constraint.T], + # [None, self.potential_constraint, None], + # ], + # format="csc", + # ) + # self.l_scheme_mixed_darcy_lu = ( + # sps.linalg.splu(l_scheme_mixed_darcy) + # ) + # + # # Reset stagnation counter + # num_neg_diff = 0 + # + # L_max = self.options.get("L_max", 1e8) + # if self.L > L_max: + # break + + # Update Bregman variables + old_flux = new_flux.copy() + old_aux_flux = new_aux_flux.copy() + old_force = new_force.copy() + old_distance = new_distance + + # TODO solve for potential and multiplier + solution_i = np.zeros_like(rhs) + solution_i[self.flux_slice] = new_flux.copy() + # TODO continue # Define performance metric status = { "converged": iter < num_iter, "number iterations": iter, "distance": new_distance, - "residual mass conservation": mass_conservation_residual, - "flux increment": flux_diff, + "mass conservation residual": error[0], + "flux increment": error[2], "distance increment": abs(new_distance - old_distance), + "convergence history": convergence_history, } return new_distance, solution_i, status @@ -855,17 +1605,18 @@ def wasserstein_distance( """ if method.lower() in ["newton", "bregman"]: - shape = mass_1.img.shape - voxel_size = mass_1.voxel_size - dim = mass_1.space_dim + # Extract grid - implicitly assume mass_2 to generate same grid + grid: darsia.Grid = darsia.generate_grid(mass_1) + + # Fetch options + options = kwargs.get("options", {}) plot_solution = kwargs.get("plot_solution", False) return_solution = kwargs.get("return_solution", False) - options = kwargs.get("options", {}) if method.lower() == "newton": - w1 = WassersteinDistanceNewton(shape, voxel_size, dim, options) + w1 = WassersteinDistanceNewton(grid, options) elif method.lower() == "bregman": - w1 = WassersteinDistanceBregman(shape, voxel_size, dim, options) + w1 = WassersteinDistanceBregman(grid, options) return w1( mass_1, mass_2, plot_solution=plot_solution, return_solution=return_solution ) diff --git a/src/darsia/utils/andersonacceleration.py b/src/darsia/utils/andersonacceleration.py index e6b7eaf7..0f98e406 100644 --- a/src/darsia/utils/andersonacceleration.py +++ b/src/darsia/utils/andersonacceleration.py @@ -1,14 +1,30 @@ -from typing import Union +"""Anderson acceleration as described by Walker and Ni in doi:10.2307/23074353.""" + +from typing import Optional, Union import numpy as np import scipy as sp class AndersonAcceleration: - """Anderson acceleration as described by Walker and Ni in doi:10.2307/23074353.""" + """Anderson acceleration for fixed-point methods.""" + + def __init__( + self, + dimension: Optional[Union[int, tuple[int]]] = None, + depth: int = 0, + restart: Optional[int] = None, + ) -> None: + """Initialize Anderson acceleration. - def __init__(self, dimension: Union[int, tuple[int]], depth: int = 0) -> None: - """Initialize Anderson acceleration.""" + Args: + dimension (int, tuple[int]): dimension of the problem. If a tuple is given, + the problem is assumed to be a tensor problem and the dimension is + calculated as the product of the tuple entries. + depth (int): depth of the acceleration. If 0, no acceleration is applied. + restart (int): restart of the acceleration. If None, no restart is applied. + + """ if isinstance(dimension, np.integer): self._dimension = dimension @@ -17,15 +33,14 @@ def __init__(self, dimension: Union[int, tuple[int]], depth: int = 0) -> None: self.tensor_shape = dimension self._dimension = int(np.prod(dimension)) self._tensor = True + elif dimension is None: + self._tensor = False + pass else: raise ValueError("Dimension not recognized.") self._depth = depth - - # Initialize arrays for iterates. - self.reset() - self._fkm1: np.ndarray = self._Fk.copy() - self._gkm1: np.ndarray = self._Gk.copy() + self._restart = restart def reset(self) -> None: """Reset Anderson acceleration.""" @@ -35,6 +50,10 @@ def reset(self) -> None: self._Gk: np.ndarray = np.zeros( (self._dimension, self._depth) ) # changes in fixed point applications + self._fkm1: np.ndarray = self._Fk.copy() + self._gkm1: np.ndarray = self._Gk.copy() + + self._inner_iteration = 0 def __call__(self, gk: np.ndarray, fk: np.ndarray, iteration: int) -> np.ndarray: """Apply Anderson acceleration. @@ -54,15 +73,18 @@ def __call__(self, gk: np.ndarray, fk: np.ndarray, iteration: int) -> np.ndarray gk = np.ravel(gk) fk = np.ravel(fk) - if iteration == 0: - self._Fk = np.zeros((self._dimension, self._depth)) # changes in increments - self._Gk = np.zeros( - (self._dimension, self._depth) - ) # changes in fixed point applications + if self._restart is not None: + self._inner_iteration = iteration % self._restart + else: + self._inner_iteration = iteration - mk = min(iteration, self._depth) + # Reset if necessary. + if self._inner_iteration == 0: + self._dimension = len(gk) + self.reset() - # Apply actual acceleration (not in the first iteration). + # Apply actual acceleration (not in the first iteration, of any restart loop). + mk = min(self._inner_iteration, self._depth) if mk > 0: # Build matrices of changes col = (iteration - 1) % self._depth diff --git a/src/darsia/utils/fv.py b/src/darsia/utils/fv.py new file mode 100644 index 00000000..ca4b74ae --- /dev/null +++ b/src/darsia/utils/fv.py @@ -0,0 +1,404 @@ +"""Finite volume utilities.""" + +import numpy as np +import scipy.sparse as sps + +import darsia + +# ! ---- Finite volume operators ---- + + +class FVDivergence: + """Finite volume divergence operator.""" + + def __init__(self, grid: darsia.Grid) -> None: + # Define sparse divergence operator, integrated over elements. + # Note: The global direction of the degrees of freedom is hereby fixed for all + # faces. In 2d, fluxes across vertical faces go from left to right, fluxes + # across horizontal faces go from bottom to top. To oppose the direction of the + # outer normal, the sign of the divergence is flipped for one side of cells for + # all faces. Analogously, in 3d. + div_shape = (grid.num_cells, grid.num_faces) + div_data = np.concatenate( + [ + grid.face_vol[d] * np.tile([1, -1], grid.num_faces_per_axis[d]) + for d in range(grid.dim) + ] + ) + div_row = np.concatenate( + [np.ravel(grid.connectivity[grid.faces[d]]) for d in range(grid.dim)] + ) + div_col = np.repeat(np.arange(grid.num_faces, dtype=int), 2) + div = sps.csc_matrix( + (div_data, (div_row, div_col)), + shape=div_shape, + ) + + # Cache + self.mat = div + + +class FVMass: + """Finite volume mass matrix. + + The mass matrix can be formulated for cell and face quantities. For cell + quantities, the mass matrix is diagonal and has the volume of the cells on the + diagonal. For face quantities, the mass matrix is diagonal and has half the cell + volumes on the diagonal, taking into accout itegration over te faces as lower + dimensional entities, and the double occurrence of the faces in the integration + over the cells. + + """ + + def __init__( + self, grid: darsia.Grid, mode: str = "cells", lumping: bool = True + ) -> None: + # Define sparse mass matrix on cells: flat_mass -> flat_mass + if mode == "cells": + mass_matrix = sps.diags( + np.prod(grid.voxel_size) * np.ones(grid.num_cells, dtype=float) + ) + elif mode == "faces": + if lumping: + mass_matrix = 0.5 * sps.diags( + np.prod(grid.voxel_size) * np.ones(grid.num_faces, dtype=float) + ) + else: + raise NotImplementedError + + # Define true RT0 mass matrix on faces: flat fluxes -> flat fluxes + num_inner_cells_with_vertical_faces = len( + grid.inner_cells_with_vertical_faces + ) + num_inner_cells_with_horizontal_faces = len( + grid.inner_cells_with_horizontal_faces + ) + mass_matrix_shape = (grid.num_faces, grid.num_faces) + mass_matrix_data = np.prod(grid.voxel_size) * np.concatenate( + ( + 2 / 3 * np.ones(grid.num_faces, dtype=float), # all faces + 1 + / 6 + * np.ones( + num_inner_cells_with_vertical_faces, dtype=float + ), # left faces + 1 + / 6 + * np.ones( + num_inner_cells_with_vertical_faces, dtype=float + ), # right faces + 1 + / 6 + * np.ones( + num_inner_cells_with_horizontal_faces, dtype=float + ), # top faces + 1 + / 6 + * np.ones( + num_inner_cells_with_horizontal_faces, dtype=float + ), # bottom faces + ) + ) + mass_matrix_row = np.concatenate( + ( + np.arange(grid.num_faces, dtype=int), + grid.connectivity_cell_to_vertical_face[ + grid.inner_cells_with_vertical_faces, 0 + ], + grid.connectivity_cell_to_vertical_face[ + grid.inner_cells_with_vertical_faces, 1 + ], + grid.connectivity_cell_to_horizontal_face[ + grid.inner_cells_with_horizontal_faces, 0 + ], + grid.connectivity_cell_to_horizontal_face[ + grid.inner_cells_with_horizontal_faces, 1 + ], + ) + ) + mass_matrix_col = np.concatenate( + ( + np.arange(grid.num_faces, dtype=int), + grid.connectivity_cell_to_vertical_face[ + grid.inner_cells_with_vertical_faces, 1 + ], + grid.connectivity_cell_to_vertical_face[ + grid.inner_cells_with_vertical_faces, 0 + ], + grid.connectivity_cell_to_horizontal_face[ + grid.inner_cells_with_horizontal_faces, 1 + ], + grid.connectivity_cell_to_horizontal_face[ + grid.inner_cells_with_horizontal_faces, 0 + ], + ) + ) + + # Define mass matrix in faces + mass_matrix = sps.csc_matrix( + ( + mass_matrix_data, + (mass_matrix_row, mass_matrix_col), + ), + shape=mass_matrix_shape, + ) + + # Cache + self.mat = mass_matrix + + +class FVTangentialFaceReconstruction: + """Projection of normal fluxes on grid onto tangential components. + + The tangential components are defined as the components of the fluxes that are + orthogonal to the normal of the face. The tangential components are determined + through averaging the fluxes on the faces that are orthogonal to the face of + interest. + + The resulting tangential flux has co-dimension 1, i.e. it is a scalar quantity in + 2d and a 2-valued vector quantity in 3d. + + """ + + def __init__(self, grid: darsia.Grid) -> None: + """Initialize the average operator. + + Args: + grid (darsia.Grid): grid + + """ + + # Operator for averaging fluxes on orthogonal, neighboring faces + shape = (grid.num_faces, grid.num_faces) + + # Each interior inner face has four neighboring faces with normal direction + # and all oriented in the same direction. In three dimensions, two such normal + # directions exist. For outer inner faces, there are only two such neighboring + # faces. + data = 0.25 * np.ones(4 * grid.num_faces, dtype=float) + + # The rows correspond to the faces for which the tangential fluxes (later to be + # tiled for addressing the right amount of vectorial components). + rows = np.concatenate([np.repeat(grid.faces[d], 4) for d in range(grid.dim)]) + + # The columns correspond to the (orthogonal) faces contributing to the average + # of the tangential fluxes. The main idea is for each face to follow the + # connectivity. + cols = [ + np.concatenate( + [ + np.ravel( + grid.reverse_connectivity[ + d_perp, + np.ravel(grid.connectivity[grid.faces[d]]), + ] + ) + for d in range(grid.dim) + for d_perp in [np.delete(range(grid.dim), d)[i]] + ] + ) + for i in range(grid.dim - 1) + ] + + # Construct and cache the sparse projection matrix, need to extract those faces + # which are not inner faces. + self.mat = [ + sps.csc_matrix( + ( + data[col != -1], + (rows[col != -1], col[col != -1]), + ), + shape=shape, + ) + for col in cols + ] + + # Cache some informatio + self.num_tangential_directions = grid.dim - 1 + self.grid = grid + + def __call__(self, normal_flux: np.ndarray, concatenate: bool = True) -> np.ndarray: + """Apply the operator to the normal fluxes. + + Args: + normal_flux (np.ndarray): normal fluxes + concatenate (bool, optional): whether to concatenate the tangential fluxes + + Returns: + np.ndarray: tangential fluxes + + """ + # Apply the operator to the normal fluxes + tangential_flux = [ + self.mat[d].dot(normal_flux) for d in range(self.num_tangential_directions) + ] + if concatenate: + tangential_flux = np.concatenate(tangential_flux, axis=0) + + return tangential_flux + + +class FVFullFaceReconstruction: + def __init__(self, grid: darsia.Grid) -> None: + self.grid = grid + self.tangential_reconstruction = FVTangentialFaceReconstruction(grid) + + def __call__(self, normal_flux: np.ndarray) -> np.ndarray: + """Reconstruct the full fluxes from the normal and tangential fluxes. + + Args: + normal_flux (np.ndarray): normal fluxes + + Returns: + np.ndarray: full fluxes + + """ + # Apply the operator to the normal fluxes + tangential_fluxes = self.tangential_reconstruction(normal_flux, False) + + # Reconstruct the full fluxes + dim = self.grid.dim + full_flux = np.zeros((self.grid.num_faces, dim), dtype=float) + for d in range(dim): + full_flux[self.grid.faces[d], d] = normal_flux[self.grid.faces[d]] + for i, d_perp in enumerate(np.delete(range(dim), d)): + full_flux[self.grid.faces[d], d_perp] = tangential_fluxes[i][ + self.grid.faces[d] + ] + + return full_flux + + +# ! ---- Finite volume projection operators ---- + + +def face_to_cell(grid: darsia.Grid, flat_flux: np.ndarray) -> np.ndarray: + """Reconstruct the vector fluxes on the cells from normal fluxes on the faces. + + Use the Raviart-Thomas reconstruction of the fluxes on the cells from the fluxes + on the faces, and use arithmetic averaging of the fluxes on the faces, + equivalent with the L2 projection of the fluxes on the faces to the fluxes on + the cells. + + Matrix-free implementation. + + Args: + grid (darsia.Grid): grid + flat_flux (np.ndarray): flat fluxes (normal fluxes on the faces) + + Returns: + np.ndarray: cell-based vectorial fluxes + + """ + # TODO revert order of indices + cell_flux = np.zeros((*grid.shape, grid.dim), dtype=float) + + if grid.dim >= 1: + cell_flux[:-1, ..., 0] += 0.5 * flat_flux[grid.faces[0]].reshape( + grid.faces_shape[0], order="F" + ) + cell_flux[1:, ..., 0] += 0.5 * flat_flux[grid.faces[0]].reshape( + grid.faces_shape[0], order="F" + ) + if grid.dim >= 2: + cell_flux[:, :-1, ..., 1] += 0.5 * flat_flux[grid.faces[1]].reshape( + grid.faces_shape[1], order="F" + ) + cell_flux[:, 1:, ..., 1] += 0.5 * flat_flux[grid.faces[1]].reshape( + grid.faces_shape[1], order="F" + ) + if grid.dim >= 3: + cell_flux[:, :, :-1, ..., 2] += 0.5 * flat_flux[grid.faces[2]].reshape( + grid.faces_shape[2], order="F" + ) + cell_flux[:, :, 1:, ..., 2] += 0.5 * flat_flux[grid.faces[2]].reshape( + grid.faces_shape[2], order="F" + ) + if grid.dim > 3: + raise NotImplementedError(f"Dimension {grid.dim} not supported.") + + return cell_flux + + +def cell_to_face(grid: darsia.Grid, cell_qty: np.ndarray, mode: str) -> np.ndarray: + """Project scalar cell quantity to scalar face quantity. + + Allow for arithmetic or harmonic averaging of the cell quantity to the faces. In + the harmonic case, the averaging is regularized to avoid division by zero. + Matrix-free implementation. + + Args: + grid (darsia.Grid): grid + cell_qty (np.ndarray): scalar-valued cell-based quantity + mode (str): mode of projection, either "arithmetic" or "harmonic" + (averaging) + + Returns: + np.ndarray: face-based quantity + + """ + + # NOTE: No impact of Grid here, so far! Everything is implicit. This should/could + # change. In particular when switching to 3d! + + raise NotImplementedError + + # Determine the fluxes on the faces + if mode == "arithmetic": + # Employ arithmetic averaging + horizontal_face_qty = 0.5 * (cell_qty[:, :-1] + cell_qty[:, 1:]) + vertical_face_qty = 0.5 * (cell_qty[:-1, :] + cell_qty[1:, :]) + elif mode == "harmonic": + # Employ harmonic averaging + arithmetic_avg_horizontal = 0.5 * (cell_qty[:, :-1] + cell_qty[:, 1:]) + arithmetic_avg_vertical = 0.5 * (cell_qty[:-1, :] + cell_qty[1:, :]) + # Regularize to avoid division by zero + regularization = 1e-10 + arithmetic_avg_horizontal = ( + arithmetic_avg_horizontal + + (2 * np.sign(arithmetic_avg_horizontal) + 1) * regularization + ) + arithmetic_avg_vertical = ( + 0.5 * arithmetic_avg_vertical + + (2 * np.sign(arithmetic_avg_vertical) + 1) * regularization + ) + product_horizontal = np.multiply(cell_qty[:, :-1], cell_qty[:, 1:]) + product_vertical = np.multiply(cell_qty[:-1, :], cell_qty[1:, :]) + + # Determine the harmonic average + horizontal_face_qty = product_horizontal / arithmetic_avg_horizontal + vertical_face_qty = product_vertical / arithmetic_avg_vertical + else: + raise ValueError(f"Mode {mode} not supported.") + + # Reshape the fluxes - hardcoding the connectivity here + face_qty = np.concatenate([horizontal_face_qty.ravel(), vertical_face_qty.ravel()]) + + return face_qty + + +# NOTE: Currently not in use. TODO rm? +# def face_restriction(self, cell_flux: np.ndarray) -> np.ndarray: +# """Restrict vector-valued fluxes on cells to normal components on faces. +# +# Matrix-free implementation. The fluxes on the faces are determined by +# arithmetic averaging of the fluxes on the cells in the direction of the normal +# of the face. +# +# Args: +# cell_flux (np.ndarray): cell-based fluxes +# +# Returns: +# np.ndarray: face-based fluxes +# +# """ +# # Determine the fluxes on the faces through arithmetic averaging +# horizontal_fluxes = 0.5 * (cell_flux[:, :-1, 0] + cell_flux[:, 1:, 0]) +# vertical_fluxes = 0.5 * (cell_flux[:-1, :, 1] + cell_flux[1:, :, 1]) +# +# # Reshape the fluxes +# flat_flux = np.concatenate( +# [horizontal_fluxes.ravel(), vertical_fluxes.ravel()], axis=0 +# ) +# +# return flat_flux diff --git a/src/darsia/utils/grid.py b/src/darsia/utils/grid.py new file mode 100644 index 00000000..6a921bcc --- /dev/null +++ b/src/darsia/utils/grid.py @@ -0,0 +1,198 @@ +"""Grid utilities for tensor grids.""" + +from typing import Union + +import numpy as np + +import darsia + +# TODO make nested lits to arrays for faster access. + + +class Grid: + """Tensor grid. + + The current implmentatio does nt take into accout true boundary faces assuming + the boundar values are not of interest. + + """ + + def __init__(self, shape: tuple, voxel_size: Union[float, list] = 1.0): + """Initialize grid.""" + + # Cache grid info + self.dim = len(shape) + """int: Number of dimensions.""" + + self.shape = shape + """tuple: Shape of grid, using matrix/tensor indexing.""" + + self.voxel_size = ( + np.array(voxel_size) + if isinstance(voxel_size, list) + else voxel_size * np.ones(self.dim) + ) + """np.ndarray: Size of voxels in each dimension.""" + + self.face_vol = [ + np.prod(self.voxel_size[np.delete(np.arange(self.dim), d)]) + for d in range(self.dim) + ] + """list: Volume of faces in each dimension.""" + + assert len(self.voxel_size) == self.dim + + # Define cell and face numbering + self._setup() + + def _setup(self) -> None: + """Define cell and face numbering.""" + + # ! ---- Grid management ---- + + # Define dimensions of the problem and indexing of cells, from here one start + # counting rows from left to right, from top to bottom. + self.num_cells = np.prod(self.shape) + """int: Number of cells.""" + + self.cell_index = np.arange(self.num_cells, dtype=int).reshape( + self.shape, order="F" + ) + """np.ndarray: cell indices, following the matrix indexing convention.""" + + # Determine number of inner faces in each axis + self.faces_shape = [ + np.array(self.shape) - np.eye(self.dim, dtype=int)[d] + for d in range(self.dim) + ] + """list: Shape of inner faces in each axis.""" + + self.num_faces_per_axis = [np.prod(s) for s in self.faces_shape] + """list: Number of inner faces in each axis.""" + + self.num_faces = np.sum(self.num_faces_per_axis) + """int: Number of faces.""" + + # Define indexing and ordering of inner faces. Horizontal -> vertical -> depth. + # TODO replace with slices + self.faces = [ + sum(self.num_faces_per_axis[:d]) + + np.arange(self.num_faces_per_axis[d], dtype=int) + for d in range(self.dim) + ] + """list: Indices of inner faces in each axis.""" + + self.face_index = [ + self.faces[d].reshape(self.faces_shape[d], order="F") + for d in range(self.dim) + ] + """list: Indices of inner faces in each axis, using matrix indexing.""" + + # Identify inner faces (full cube) + self.interior_faces = [] + """list: Indices of interior faces.""" + + if self.dim == 1: + self.interior_faces = [ + np.ravel(self.face_index[0][1:-1], "F"), + ] + elif self.dim == 2: + self.interior_faces = [ + np.ravel(self.face_index[0][:, 1:-1], "F"), + np.ravel(self.face_index[1][1:-1, :], "F"), + ] + elif self.dim == 3: + self.interior_faces = [ + np.ravel(self.face_index[0][:, 1:-1, 1:-1], "F"), + np.ravel(self.face_index[1][1:-1, :, 1:-1], "F"), + np.ravel(self.face_index[2][1:-1, 1:-1, :], "F"), + ] + else: + raise NotImplementedError(f"Grid of dimension {self.dim} not implemented.") + + # Identify all faces on the outer boundary of the grid. Need to use hardcoded + # knowledge of the orientation of axes and grid indexing. + self.exterior_faces = [ + np.sort(np.array(list(set(self.faces[d]) - set(self.interior_faces[d])))) + for d in range(self.dim) + ] + + # ! ---- Connectivity ---- + + self.connectivity = np.zeros((self.num_faces, 2), dtype=int) + """np.ndarray: Connectivity (and direction) of faces to cells.""" + + if self.dim >= 1: + self.connectivity[self.faces[0], 0] = np.ravel( + self.cell_index[:-1, ...], "F" + ) + self.connectivity[self.faces[0], 1] = np.ravel( + self.cell_index[1:, ...], "F" + ) + if self.dim >= 2: + self.connectivity[self.faces[1], 0] = np.ravel( + self.cell_index[:, :-1, ...], "F" + ) + self.connectivity[self.faces[1], 1] = np.ravel( + self.cell_index[:, 1:, ...], "F" + ) + if self.dim >= 3: + self.connectivity[self.faces[2], 0] = np.ravel( + self.cell_index[:, :, :-1, ...], "F" + ) + self.connectivity[self.faces[2], 1] = np.ravel( + self.cell_index[:, :, 1:, ...], "F" + ) + if self.dim > 3: + raise NotImplementedError(f"Grid of dimension {self.dim} not implemented.") + + self.reverse_connectivity = -np.ones((self.dim, self.num_cells, 2), dtype=int) + """np.ndarray: Reverse connectivity (and direction) of cells to faces.""" + + # NOTE: The first components addresses the normal direction of the face, + # the second the cell, the third the relative position of the face in the + # cell. + + if self.dim >= 1: + self.reverse_connectivity[ + 0, np.ravel(self.cell_index[1:, ...], "F"), 0 + ] = self.faces[0] + self.reverse_connectivity[ + 0, np.ravel(self.cell_index[:-1, ...], "F"), 1 + ] = self.faces[0] + + if self.dim >= 2: + self.reverse_connectivity[ + 1, np.ravel(self.cell_index[:, 1:, ...], "F"), 0 + ] = self.faces[1] + self.reverse_connectivity[ + 1, np.ravel(self.cell_index[:, :-1, ...], "F"), 1 + ] = self.faces[1] + + if self.dim >= 3: + self.reverse_connectivity[ + 2, np.ravel(self.cell_index[:, :, 1:, ...], "F"), 0 + ] = self.faces[2] + self.reverse_connectivity[ + 2, np.ravel(self.cell_index[:, :, :-1, ...], "F"), 1 + ] = self.faces[2] + + # Info about inner cells + # TODO rm? + self.inner_cells_with_inner_faces = ( + [] + [np.ravel(self.cell_index[1:-1, ...], "F")] + if self.dim >= 1 + else [] + [np.ravel(self.cell_index[:, 1:-1, ...], "F")] + if self.dim >= 2 + else [] + [np.ravel(self.cell_index[:, :, 1:-1, ...], "F")] + if self.dim >= 3 + else [] + ) + """list: Indices of inner cells with inner faces.""" + + +def generate_grid(image: darsia.Image) -> Grid: + """Get grid object.""" + shape = image.num_voxels + voxel_size = image.voxel_size + return Grid(shape, voxel_size) diff --git a/tests/unit/test_fv.py b/tests/unit/test_fv.py new file mode 100644 index 00000000..3f5f8e29 --- /dev/null +++ b/tests/unit/test_fv.py @@ -0,0 +1,278 @@ +"""Unit tests for finite volume utilities.""" + +import numpy as np + +import darsia + + +def test_divergence_2d(): + # Create divergence matrix + grid = darsia.Grid(shape=(4, 5), voxel_size=[0.5, 0.25]) + divergence = darsia.FVDivergence(grid).mat.todense() + + # Check shape + assert np.allclose(divergence.shape, (grid.num_cells, grid.num_faces)) + + # Check values in corner cells + assert np.allclose(np.nonzero(divergence[0])[1], [0, 15]) + assert np.allclose(divergence[0, np.array([0, 15])], [0.25, 0.5]) + + assert np.allclose(np.nonzero(divergence[4])[1], [3, 15, 19]) + assert np.allclose(divergence[4, np.array([3, 15, 19])], [0.25, -0.5, 0.5]) + + assert np.allclose(np.nonzero(divergence[16])[1], [12, 27]) + assert np.allclose(divergence[16, np.array([12, 27])], [0.25, -0.5]) + + assert np.allclose(np.nonzero(divergence[19])[1], [14, 30]) + assert np.allclose(divergence[19, np.array([14, 30])], [-0.25, -0.5]) + + # Check value for interior cell + assert np.allclose(np.nonzero(divergence[6])[1], [4, 5, 17, 21]) + assert np.allclose( + divergence[6, np.array([4, 5, 17, 21])], [-0.25, 0.25, -0.5, 0.5] + ) + + +def test_divergence_3d(): + # Create divergence matrix + grid = darsia.Grid(shape=(3, 4, 5), voxel_size=[0.5, 0.25, 2]) + divergence = darsia.FVDivergence(grid).mat.todense() + + # Check shape + assert np.allclose(divergence.shape, (grid.num_cells, grid.num_faces)) + + # Check values in corner cells + assert np.allclose(np.nonzero(divergence[0])[1], [0, 40, 85]) + assert np.allclose(divergence[0, np.array([0, 40, 85])], [0.5, 1, 0.125]) + + assert np.allclose(np.nonzero(divergence[11])[1], [7, 48, 96]) + assert np.allclose(divergence[11, np.array([7, 48, 96])], [-0.5, -1, 0.125]) + + assert np.allclose(np.nonzero(divergence[59])[1], [39, 84, 132]) + assert np.allclose(divergence[59, np.array([39, 84, 132])], [-0.5, -1, -0.125]) + + # Check value for interior cell + assert np.allclose(np.nonzero(divergence[16])[1], [10, 11, 50, 53, 89, 101]) + assert np.allclose( + divergence[16, np.array([10, 11, 50, 53, 89, 101])], + [-0.5, 0.5, -1, 1, -0.125, 0.125], + ) + + +def test_mass_2d(): + grid = darsia.Grid(shape=(4, 5), voxel_size=[0.5, 0.25]) + mass = darsia.FVMass(grid).mat.todense() + + # Check shape + assert np.allclose(mass.shape, (grid.num_cells, grid.num_cells)) + + # Check diagonal structure + assert np.linalg.norm(mass - np.diag(np.diag(mass))) < 1e-10 + + # Check diagonal values + assert len(np.unique(np.diag(mass))) == 1 + assert np.isclose(mass[0, 0], 0.125) + + +def test_mass_3d(): + grid = darsia.Grid(shape=(3, 4, 5), voxel_size=[0.5, 0.25, 2]) + mass = darsia.FVMass(grid).mat.todense() + + # Check shape + assert np.allclose(mass.shape, (grid.num_cells, grid.num_cells)) + + # Check diagonal structure + assert np.linalg.norm(mass - np.diag(np.diag(mass))) < 1e-10 + + # Check diagonal values + assert len(np.unique(np.diag(mass))) == 1 + assert np.isclose(mass[0, 0], 0.25) + + +def test_mass_face_2d(): + grid = darsia.Grid(shape=(4, 5), voxel_size=[0.5, 0.25]) + mass = darsia.FVMass(grid, mode="faces", lumping=True).mat.todense() + + # Check shape + assert np.allclose(mass.shape, (grid.num_faces, grid.num_faces)) + + # Check diagonal structure + assert np.linalg.norm(mass - np.diag(np.diag(mass))) < 1e-10 + + # Check diagonal values + assert len(np.unique(np.diag(mass))) == 1 + assert np.isclose(mass[0, 0], 0.5 * 0.125) + + +def test_mass_face_3d(): + grid = darsia.Grid(shape=(3, 4, 5), voxel_size=[0.5, 0.25, 2]) + mass = darsia.FVMass(grid, mode="faces", lumping=True).mat.todense() + + # Check shape + assert np.allclose(mass.shape, (grid.num_faces, grid.num_faces)) + + # Check diagonal structure + assert np.linalg.norm(mass - np.diag(np.diag(mass))) < 1e-10 + + # Check diagonal values + assert len(np.unique(np.diag(mass))) == 1 + assert np.isclose(mass[0, 0], 0.5 * 0.25) + + +def test_tangential_reconstruction_2d_1(): + grid = darsia.Grid(shape=(2, 3), voxel_size=[0.5, 0.25]) + tangential_reconstruction = ( + darsia.FVTangentialFaceReconstruction(grid).mat[0].todense() + ) + + # Check shape + assert np.allclose( + tangential_reconstruction.shape, (grid.num_faces, grid.num_faces) + ) + + # Check values - first for exterior faces + assert np.allclose(tangential_reconstruction[0], [0, 0, 0, 0.25, 0.25, 0, 0]) + assert np.allclose(tangential_reconstruction[4], [0.25, 0.25, 0, 0, 0, 0, 0]) + + # Check values - then for interior faces + assert np.allclose(tangential_reconstruction[1], [0, 0, 0, 0.25, 0.25, 0.25, 0.25]) + + +def test_tangential_reconstruction_2d_2(): + grid = darsia.Grid(shape=(3, 4), voxel_size=[0.5, 0.25]) + tangential_reconstruction_dense = ( + darsia.FVTangentialFaceReconstruction(grid).mat[0].todense() + ) + + # Check shape + assert np.allclose( + tangential_reconstruction_dense.shape, (grid.num_faces, grid.num_faces) + ) + + # Check values - first for exterior faces + assert np.allclose(np.nonzero(tangential_reconstruction_dense[0])[1], [8, 9]) + assert np.allclose(np.nonzero(tangential_reconstruction_dense[7])[1], [15, 16]) + + # Check values - then for interior faces + assert np.allclose( + np.nonzero(tangential_reconstruction_dense[2])[1], [8, 9, 11, 12] + ) + assert np.allclose(np.nonzero(tangential_reconstruction_dense[15])[1], [4, 5, 6, 7]) + + # Apply once and prove values + tangential_reconstruction = darsia.FVTangentialFaceReconstruction(grid) + normal_flux = np.arange(grid.num_faces) + tangential_flux = tangential_reconstruction(normal_flux) + assert np.allclose( + tangential_flux[np.array([0, 1, 4, 8, 12, 16])], [4.25, 4.75, 13, 0.5, 3.5, 3] + ) + + +def test_full_reconstruction_2d_2(): + grid = darsia.Grid(shape=(3, 4), voxel_size=[0.5, 0.25]) + + # Apply once and prove values + tangential_reconstruction = darsia.FVFullFaceReconstruction(grid) + normal_flux = np.arange(grid.num_faces) + full_flux = tangential_reconstruction(normal_flux) + + # Check shape + assert np.allclose(full_flux.shape, (grid.num_faces, 2)) + + # Check values + assert np.allclose(full_flux[0], [0, 4.25]) + assert np.allclose(full_flux[1], [1, 4.75]) + assert np.allclose(full_flux[4], [4, 13]) + assert np.allclose(full_flux[8], [0.5, 8]) + assert np.allclose(full_flux[12], [3.5, 12]) + assert np.allclose(full_flux[16], [3, 16]) + + +def test_tangential_reconstruction_3d(): + grid = darsia.Grid(shape=(3, 3, 3), voxel_size=[0.5, 0.25, 2]) + tangential_reconstruction = darsia.FVTangentialFaceReconstruction(grid) + + # Apply once and probe values + normal_flux = np.arange(grid.num_faces) + tangential_flux = tangential_reconstruction(normal_flux, concatenate=False) + + # Corner block + assert np.allclose( + [tangential_flux[0][0], tangential_flux[1][0]], [(18 + 19) / 4, (36 + 37) / 4] + ) + assert np.allclose( + [tangential_flux[0][18], tangential_flux[1][18]], [(0 + 2) / 4, (36 + 39) / 4] + ) + assert np.allclose( + [tangential_flux[0][36], tangential_flux[1][36]], [(0 + 6) / 4, (18 + 24) / 4] + ) + + # Center block + assert np.allclose( + [tangential_flux[0][8], tangential_flux[1][8]], + [(24 + 25 + 27 + 28) / 4, (39 + 40 + 48 + 49) / 4], + ) + assert np.allclose( + [tangential_flux[0][25], tangential_flux[1][25]], + [(6 + 7 + 8 + 9) / 4, (37 + 40 + 46 + 49) / 4], + ) + assert np.allclose( + [tangential_flux[0][40], tangential_flux[1][40]], + [(2 + 3 + 8 + 9) / 4, (19 + 22 + 25 + 28) / 4], + ) + + +def test_full_reconstruction_3d(): + grid = darsia.Grid(shape=(3, 3, 3), voxel_size=[0.5, 0.25, 2]) + full_reconstuction = darsia.FVFullFaceReconstruction(grid) + + # Apply once and probe values + normal_flux = np.arange(grid.num_faces) + full_flux = full_reconstuction(normal_flux) + + # Corner block + assert np.allclose(full_flux[0], [0, (18 + 19) / 4, (36 + 37) / 4]) + assert np.allclose(full_flux[18], [(0 + 2) / 4, 18, (36 + 39) / 4]) + assert np.allclose(full_flux[36], [(0 + 6) / 4, (18 + 24) / 4, 36]) + + # Center block + assert np.allclose( + full_flux[8], [8, (24 + 25 + 27 + 28) / 4, (39 + 40 + 48 + 49) / 4] + ) + assert np.allclose( + full_flux[25], [(6 + 7 + 8 + 9) / 4, 25, (37 + 40 + 46 + 49) / 4] + ) + assert np.allclose( + full_flux[40], [(2 + 3 + 8 + 9) / 4, (19 + 22 + 25 + 28) / 4, 40] + ) + + +def test_face_to_cell_2d(): + grid = darsia.Grid(shape=(3, 4), voxel_size=[0.5, 0.25]) + num_faces = grid.num_faces + flat_flux = np.arange(num_faces) + cell_flux = darsia.face_to_cell(grid, flat_flux) + + # Check shape + assert np.allclose(cell_flux.shape, (*grid.shape, grid.dim)) + + # Check values + print(cell_flux) + assert np.allclose(cell_flux[0, 0], [0, 4]) + assert np.allclose(cell_flux[2, 3], [3.5, 8]) + assert np.allclose(cell_flux[1, 1], [2.5, 10.5]) + + +def test_face_to_cell_3d(): + grid = darsia.Grid(shape=(3, 4, 5), voxel_size=[0.5, 0.25, 2]) + num_faces = grid.num_faces + flat_flux = np.arange(num_faces) + cell_flux = darsia.face_to_cell(grid, flat_flux) + + # Check shape + assert np.allclose(cell_flux.shape, (*grid.shape, grid.dim)) + + # Check values + assert np.allclose(cell_flux[0, 0, 0], [0, 20, 42.5]) + assert np.allclose(cell_flux[2, 3, 4], [19.5, 42, 66]) + assert np.allclose(cell_flux[1, 1, 1], [10.5, 51.5, 95]) diff --git a/tests/unit/test_grid.py b/tests/unit/test_grid.py new file mode 100644 index 00000000..643b5596 --- /dev/null +++ b/tests/unit/test_grid.py @@ -0,0 +1,275 @@ +"""Unit tests for the grid module.""" + +import numpy as np +import pytest + +import darsia + + +def test_grid_2d(): + """Test basic indexing for 2d tensor grids.""" + + # Fetch grid + grid = darsia.Grid(shape=(3, 4), voxel_size=[0.5, 0.25]) + + # Check basic attributes + assert np.allclose(grid.shape, (3, 4)) + assert grid.dim == 2 + assert np.allclose(grid.voxel_size, [0.5, 0.25]) + + # Probe cell numbering + assert grid.num_cells == 12 + assert grid.cell_index[0, 0] == 0 + assert grid.cell_index[2, 0] == 2 + assert grid.cell_index[0, 3] == 9 + assert grid.cell_index[2, 3] == 11 + + # Check face volumes + assert np.allclose(grid.face_vol, [0.25, 0.5]) + + # Check face shape + assert np.allclose(grid.faces_shape[0], (2, 4)) + assert np.allclose(grid.faces_shape[1], (3, 3)) + + # Check face numbering + assert grid.num_faces_per_axis[0] == 8 + assert grid.num_faces_per_axis[1] == 9 + assert grid.num_faces == 17 + + # Check indexing of faces in flat format + assert np.allclose(grid.faces[0], np.arange(0, 8)) + assert np.allclose(grid.faces[1], np.arange(8, 17)) + + # Check indexing of faces in 2d format + assert np.allclose(grid.face_index[0][:, 0], np.arange(0, 2)) + assert np.allclose(grid.face_index[0][:, 1], np.arange(2, 4)) + assert np.allclose(grid.face_index[0][:, 2], np.arange(4, 6)) + assert np.allclose(grid.face_index[1][:, 0], np.arange(8, 11)) + assert np.allclose(grid.face_index[1][:, 1], np.arange(11, 14)) + assert np.allclose(grid.face_index[1][:, 2], np.arange(14, 17)) + + # Check identification of interior inner faces + assert np.allclose(grid.interior_faces[0], [2, 3, 4, 5]) + assert np.allclose(grid.interior_faces[1], [9, 12, 15]) + + # Check identification of exterior inner faces + assert np.allclose(grid.exterior_faces[0], [0, 1, 6, 7]) + assert np.allclose(grid.exterior_faces[1], [8, 10, 11, 13, 14, 16]) + + +def test_grid_connectivity_2d(): + """Test connectivity for 2d tensor grids.""" + + # Fetch grid + grid = darsia.Grid(shape=(3, 4)) + + # Check connectivity: face to cells with positive orientation + # Horizontal faces + assert np.allclose(grid.connectivity[0], [0, 1]) + assert np.allclose(grid.connectivity[1], [1, 2]) + assert np.allclose(grid.connectivity[2], [3, 4]) + assert np.allclose(grid.connectivity[3], [4, 5]) + assert np.allclose(grid.connectivity[4], [6, 7]) + assert np.allclose(grid.connectivity[5], [7, 8]) + assert np.allclose(grid.connectivity[6], [9, 10]) + assert np.allclose(grid.connectivity[7], [10, 11]) + # Vertical faces + assert np.allclose(grid.connectivity[8], [0, 3]) + assert np.allclose(grid.connectivity[9], [1, 4]) + assert np.allclose(grid.connectivity[10], [2, 5]) + assert np.allclose(grid.connectivity[11], [3, 6]) + assert np.allclose(grid.connectivity[12], [4, 7]) + assert np.allclose(grid.connectivity[13], [5, 8]) + assert np.allclose(grid.connectivity[14], [6, 9]) + assert np.allclose(grid.connectivity[15], [7, 10]) + assert np.allclose(grid.connectivity[16], [8, 11]) + + # Check reverse connectivity: cell to faces with positive orientation + # For corner cells + assert np.allclose(grid.reverse_connectivity[0, 0], [-1, 0]) + assert np.allclose(grid.reverse_connectivity[1, 0], [-1, 8]) + assert np.allclose(grid.reverse_connectivity[0, 2], [1, -1]) + assert np.allclose(grid.reverse_connectivity[1, 2], [-1, 10]) + assert np.allclose(grid.reverse_connectivity[0, 9], [-1, 6]) + assert np.allclose(grid.reverse_connectivity[1, 9], [14, -1]) + assert np.allclose(grid.reverse_connectivity[0, 11], [7, -1]) + assert np.allclose(grid.reverse_connectivity[1, 11], [16, -1]) + + # For interior cells + assert np.allclose(grid.reverse_connectivity[0, 4], [2, 3]) + assert np.allclose(grid.reverse_connectivity[1, 4], [9, 12]) + + +def test_grid_3d(): + """Test basic indexing for 2d tensor grids.""" + + # Fetch grid + grid = darsia.Grid(shape=(3, 4, 5), voxel_size=[0.5, 0.25, 2]) + + # Check basic attributes + assert np.allclose(grid.shape, (3, 4, 5)) + assert grid.dim == 3 + assert np.allclose(grid.voxel_size, [0.5, 0.25, 2]) + + # Probe cell numbering + assert grid.num_cells == 60 + assert grid.cell_index[0, 0, 0] == 0 + assert grid.cell_index[2, 0, 0] == 2 + assert grid.cell_index[0, 3, 0] == 9 + assert grid.cell_index[2, 3, 0] == 11 + assert grid.cell_index[0, 0, 4] == 48 + assert grid.cell_index[2, 0, 4] == 50 + assert grid.cell_index[0, 3, 4] == 57 + assert grid.cell_index[2, 3, 4] == 59 + + # Check face volumes + assert np.allclose(grid.face_vol, [0.5, 1, 0.125]) + + # Check face shape + assert np.allclose(grid.faces_shape[0], (2, 4, 5)) + assert np.allclose(grid.faces_shape[1], (3, 3, 5)) + assert np.allclose(grid.faces_shape[2], (3, 4, 4)) + + # Check face numbering + assert grid.num_faces_per_axis[0] == 40 + assert grid.num_faces_per_axis[1] == 45 + assert grid.num_faces_per_axis[2] == 48 + assert grid.num_faces == 40 + 45 + 48 + + # Check indexing of faces in flat format + assert np.allclose(grid.faces[0], np.arange(0, 40)) + assert np.allclose(grid.faces[1], np.arange(40, 40 + 45)) + assert np.allclose(grid.faces[2], np.arange(40 + 45, 40 + 45 + 48)) + + # Check indexing of faces in 2d format + assert np.allclose(grid.face_index[0][:, 0, 0], np.arange(0, 2)) + assert np.allclose(grid.face_index[0][:, 1, 0], np.arange(2, 4)) + assert np.allclose(grid.face_index[0][:, 2, 0], np.arange(4, 6)) + assert np.allclose(grid.face_index[0][:, 0, 4], np.arange(32, 34)) + assert np.allclose(grid.face_index[0][:, 1, 4], np.arange(34, 36)) + assert np.allclose(grid.face_index[0][:, 2, 4], np.arange(36, 38)) + assert np.allclose(grid.face_index[1][:, 0, 0], np.arange(40, 43)) + assert np.allclose(grid.face_index[1][:, 1, 0], np.arange(43, 46)) + assert np.allclose(grid.face_index[1][:, 2, 0], np.arange(46, 49)) + + # Check identification of interior inner faces + assert np.allclose( + grid.interior_faces[0], + [ + 10, + 11, + 12, + 13, + 18, + 19, + 20, + 21, + 26, + 27, + 28, + 29, + ], + ) + assert np.allclose(grid.interior_faces[1], [50, 53, 56, 59, 62, 65, 68, 71, 74]) + assert np.allclose(grid.interior_faces[2], [89, 92, 101, 104, 113, 116, 125, 128]) + + # Check identification of exterior inner faces + for d in range(grid.dim): + assert np.allclose( + np.sort(grid.exterior_faces[d]), + np.sort( + np.array( + list( + set(grid.faces[d].tolist()) + - set(grid.interior_faces[d].tolist()) + ) + ) + ), + ) + + +def test_grid_connectivity_3d(): + """Test connectivity for 3d tensor grids.""" + + # Fetch grid + grid = darsia.Grid(shape=(3, 4, 5), voxel_size=[0.5, 0.25, 2]) + + # Check connectivity: face to cells with positive orientation + # Horizontal faces + assert np.allclose(grid.connectivity[0], [0, 1]) + assert np.allclose(grid.connectivity[1], [1, 2]) + assert np.allclose(grid.connectivity[2], [3, 4]) + assert np.allclose(grid.connectivity[3], [4, 5]) + assert np.allclose(grid.connectivity[4], [6, 7]) + assert np.allclose(grid.connectivity[5], [7, 8]) + assert np.allclose(grid.connectivity[6], [9, 10]) + assert np.allclose(grid.connectivity[7], [10, 11]) + # Vertical faces + assert np.allclose(grid.connectivity[40], [0, 3]) + assert np.allclose(grid.connectivity[41], [1, 4]) + assert np.allclose(grid.connectivity[42], [2, 5]) + assert np.allclose(grid.connectivity[43], [3, 6]) + assert np.allclose(grid.connectivity[44], [4, 7]) + assert np.allclose(grid.connectivity[45], [5, 8]) + assert np.allclose(grid.connectivity[46], [6, 9]) + assert np.allclose(grid.connectivity[47], [7, 10]) + assert np.allclose(grid.connectivity[48], [8, 11]) + # Table faces + assert np.allclose(grid.connectivity[85], [0, 12]) + assert np.allclose(grid.connectivity[86], [1, 13]) + assert np.allclose(grid.connectivity[87], [2, 14]) + assert np.allclose(grid.connectivity[88], [3, 15]) + assert np.allclose(grid.connectivity[89], [4, 16]) + assert np.allclose(grid.connectivity[90], [5, 17]) + assert np.allclose(grid.connectivity[91], [6, 18]) + assert np.allclose(grid.connectivity[92], [7, 19]) + assert np.allclose(grid.connectivity[93], [8, 20]) + + # Check reverse connectivity: cell to faces with positive orientation + # For corner cells + assert np.allclose(grid.reverse_connectivity[0, 0], [-1, 0]) + assert np.allclose(grid.reverse_connectivity[1, 0], [-1, 40]) + assert np.allclose(grid.reverse_connectivity[2, 0], [-1, 85]) + assert np.allclose(grid.reverse_connectivity[0, 9], [-1, 6]) + assert np.allclose(grid.reverse_connectivity[1, 9], [46, -1]) + assert np.allclose(grid.reverse_connectivity[2, 9], [-1, 94]) + assert np.allclose(grid.reverse_connectivity[0, 57], [-1, 38]) + assert np.allclose(grid.reverse_connectivity[1, 57], [82, -1]) + assert np.allclose(grid.reverse_connectivity[2, 57], [130, -1]) + assert np.allclose(grid.reverse_connectivity[0, 59], [39, -1]) + assert np.allclose(grid.reverse_connectivity[1, 59], [84, -1]) + assert np.allclose(grid.reverse_connectivity[2, 59], [132, -1]) + + # For interior cells + assert np.allclose(grid.reverse_connectivity[0, 16], [10, 11]) + assert np.allclose(grid.reverse_connectivity[1, 16], [50, 53]) + assert np.allclose(grid.reverse_connectivity[2, 16], [89, 101]) + + +@pytest.mark.parametrize("shape", [(3, 4), (3, 4, 5)]) +def test_compatibility(shape): + """Compatibility of connectivity, reverse connectivity and interior faces. + + Make sure that inner faces really only connect to two cells which have inner faces + in all directions again. + + """ + grid = darsia.Grid(shape=shape) + assert ( + np.count_nonzero( + np.concatenate( + [ + np.ravel( + grid.reverse_connectivity[ + d_perp, + np.ravel(grid.connectivity[grid.interior_faces[d]]), + ] + ) + for d in range(grid.dim) + for d_perp in np.delete(range(grid.dim), d) + ] + ) + == -1 + ) + == 0 + ) diff --git a/tests/unit/test_wasserstein.py b/tests/unit/test_wasserstein.py new file mode 100644 index 00000000..7c2ccbe9 --- /dev/null +++ b/tests/unit/test_wasserstein.py @@ -0,0 +1,219 @@ +"""Example for Wasserstein computations moving a square to another location.""" + +import numpy as np +import pytest + +import darsia + +# ! ---- 2d version ---- + +# Coarse src image +rows = 10 +cols = rows +src_square_2d = np.zeros((rows, cols), dtype=float) +src_square_2d[2:5, 2:5] = 1 +meta_2d = {"width": 1, "height": 1, "dim": 2, "scalar": True} +src_image_2d = darsia.Image(src_square_2d, **meta_2d) + +# Coarse dst image +dst_squares_2d = np.zeros((rows, cols), dtype=float) +dst_squares_2d[1:3, 1:2] = 1 +dst_squares_2d[4:7, 7:9] = 1 +dst_image_2d = darsia.Image(dst_squares_2d, **meta_2d) + +# Rescale +shape_meta_2d = src_image_2d.shape_metadata() +geometry_2d = darsia.Geometry(**shape_meta_2d) +src_image_2d.img /= geometry_2d.integrate(src_image_2d) +dst_image_2d.img /= geometry_2d.integrate(dst_image_2d) + +# Reference value for comparison +true_distance_2d = 0.379543951823 + +# ! ---- 3d version ---- + +# Coarse src image +pages = 1 +src_square_3d = np.zeros((rows, cols, pages), dtype=float) +src_square_3d[2:5, 2:5, 0] = 1 +meta_3d = {"dimensions": [1, 1, 1], "dim": 3, "series": False, "scalar": True} +src_image_3d = darsia.Image(src_square_3d, **meta_3d) + +# Coarse dst image +dst_squares_3d = np.zeros((rows, cols, pages), dtype=float) +dst_squares_3d[1:3, 1:2, 0] = 1 +dst_squares_3d[4:7, 7:9, 0] = 1 +dst_image_3d = darsia.Image(dst_squares_3d, **meta_3d) + +# Rescale +shape_meta_3d = src_image_3d.shape_metadata() +geometry_3d = darsia.Geometry(**shape_meta_3d) +src_image_3d.img /= geometry_3d.integrate(src_image_3d) +dst_image_3d.img /= geometry_3d.integrate(dst_image_3d) + +# Reference value for comparison +true_distance_3d = 0.379543951823 + +# ! ---- Data set ---- +src_image = { + 2: src_image_2d, + 3: src_image_3d, +} + +dst_image = { + 2: dst_image_2d, + 3: dst_image_3d, +} + +true_distance = { + 2: true_distance_2d, + 3: true_distance_3d, +} + +# ! ---- Solver options ---- + +# Linearization +newton_options = { + # Scheme + "L": 1e-9, +} +bregman_std_options = { + # Scheme + "L": 1, +} +bregman_reordered_options = { + # Scheme + "L": 1, + "bregman_mode": "reordered", +} +bregman_adaptive_options = { + # Scheme + "L": 1, + "bregman_mode": "adaptive", + "bregman_update_cond": lambda iter: iter % 20 == 0, +} +linearizations = { + "newton": [newton_options], + "bregman": [ + bregman_std_options, + bregman_reordered_options, + bregman_adaptive_options, + ], +} + +# Acceleration +off_aa = { + # Nonlinear solver + "aa_depth": 0, + "aa_restart": None, +} +on_aa = { + # Nonlinear solver + "aa_depth": 5, + "aa_restart": 5, +} +accelerations = [off_aa, on_aa] + +# Linear solver +lu_options = { + # Linear solver + "linear_solver": "lu" +} +amg_options = { + "linear_solver": "amg-potential", + "linear_solver_tol": 1e-8, +} +solvers = [lu_options, amg_options] + +# General options +options = { + # Solver parameters + "regularization": 1e-16, + # Scheme + "lumping": True, + # Performance control + "num_iter": 400, + "tol_residual": 1e-10, + "tol_increment": 1e-6, + "tol_distance": 1e-10, + # Output + "verbose": False, +} + +# ! ---- Tests ---- + + +@pytest.mark.parametrize("a_key", range(len(accelerations))) +@pytest.mark.parametrize("s_key", range(len(solvers))) +@pytest.mark.parametrize("dim", [2, 3]) +def test_newton(a_key, s_key, dim): + """Test all combinations for Newton.""" + options.update(newton_options) + options.update(accelerations[a_key]) + options.update(solvers[s_key]) + distance, _, _, _, status = darsia.wasserstein_distance( + src_image[dim], + dst_image[dim], + options=options, + method="newton", + return_solution=True, + ) + assert np.isclose(distance, true_distance[dim], atol=1e-5) + assert status["converged"] + + +@pytest.mark.parametrize("a_key", range(len(accelerations))) +@pytest.mark.parametrize("s_key", range(len(solvers))) +@pytest.mark.parametrize("dim", [2, 3]) +def test_std_bregman(a_key, s_key, dim): + """Test all combinations for std Bregman.""" + options.update(bregman_std_options) + options.update(accelerations[a_key]) + options.update(solvers[s_key]) + distance, _, _, _, status = darsia.wasserstein_distance( + src_image[dim], + dst_image[dim], + options=options, + method="bregman", + return_solution=True, + ) + assert np.isclose(distance, true_distance[dim], atol=1e-2) # TODO + assert status["converged"] + + +@pytest.mark.parametrize("a_key", range(len(accelerations))) +@pytest.mark.parametrize("s_key", range(len(solvers))) +@pytest.mark.parametrize("dim", [2, 3]) +def test_reordered_bregman(a_key, s_key, dim): + """Test all combinations for reordered Bregman.""" + options.update(bregman_reordered_options) + options.update(accelerations[a_key]) + options.update(solvers[s_key]) + distance, _, _, _, status = darsia.wasserstein_distance( + src_image[dim], + dst_image[dim], + options=options, + method="bregman", + return_solution=True, + ) + assert np.isclose(distance, true_distance[dim], atol=1e-2) # TODO + assert status["converged"] + + +@pytest.mark.parametrize("a_key", range(len(accelerations))) +@pytest.mark.parametrize("s_key", range(len(solvers))) +@pytest.mark.parametrize("dim", [2, 3]) +def test_adaptive_bregman(a_key, s_key, dim): + """Test all combinations for adaptive Bregman.""" + options.update(bregman_adaptive_options) + options.update(accelerations[a_key]) + options.update(solvers[s_key]) + distance, _, _, _, status = darsia.wasserstein_distance( + src_image[dim], + dst_image[dim], + options=options, + method="bregman", + return_solution=True, + ) + assert np.isclose(distance, true_distance[dim], atol=1e-5) + assert status["converged"]