a1_mesh_and_datamap_MC.py

# %%
import numpy as np
from shapely.geometry import Polygon, Point
import sup_fun as spf
import pygimli as pg
import pygimli.meshtools as mt
from pygimli.physics import ert
from multiprocessing import Pool
import os
import time
import scipy.io

# Load optimized data
mat_data = scipy.io.loadmat('DREAM_ZS.mat')
optimized_data = mat_data['Z'][~np.all(mat_data['Z'] == 0, axis=1)]
nz = optimized_data.shape[0]
range2 = np.floor(nz / 3).astype(int)
rr = np.arange(nz - range2, nz)
optimized_par = optimized_data[rr,:-2]
# % ParRange.codes = {'aVG','nVG','cwat' ,'dwat', theta_r,
# % 'WastePar_Re.sigW1', 'WastePar_Re.sigW2' ...'WastePar_Re.sigW6'
# % 'WastePar_Re.msig','WastePar_Re.nsig'; 'WastePar_Re.C1  % 10^';
# % D, sigAT1,sigAT2,sigAT3,sigAT4,sigAT5,sigAT6;
# % feH,frel

# %%
# Define the WaterRetentionCurve class
class WaterRetentionCurve:
    def __init__(self, theta_s, theta_r, alpha, n):
        self.theta_s = theta_s  # Saturated water content
        self.theta_r = theta_r  # Residual water content
        self.alpha = alpha
        self.n = n
    
    def water_content(self, h):
        m = 1 - 1 / self.n
        return self.theta_r + (self.theta_s - self.theta_r) / (1 + (self.alpha * abs(h)) ** self.n) ** m

# Function to apply the water retention curve to a column
def apply_water_retention_curve_in_column_final(column_data, retention_curve):
    column_water_content = np.full_like(column_data, retention_curve.theta_r)  # Start with residual water content

    # Find indices where the column is saturated
    saturated_indices = np.nonzero(column_data)[0]
    column_water_content[saturated_indices] = retention_curve.theta_s

    # Process unsaturated zones
    if len(saturated_indices) > 0:
        # Below the first saturated index
        if saturated_indices[0] > 0:
            for i in range(saturated_indices[0] - 1, -1, -1):
                column_water_content[i] = retention_curve.water_content(i)
        # Between saturated indices
        for k in range(len(saturated_indices) - 1):
            for i in range(saturated_indices[k] + 1, saturated_indices[k + 1]):
                column_water_content[i] = retention_curve.water_content(i - saturated_indices[k])
        # Above the last saturated index
        if saturated_indices[-1] < len(column_data) - 1:
            for i in range(saturated_indices[-1] + 1, len(column_data)):
                column_water_content[i] = retention_curve.water_content(i - saturated_indices[-1])
    else:
        for i in range(len(column_water_content)):
            column_water_content[i] = retention_curve.water_content(i + 1)

    return column_water_content

# Function to generate irregular zones without overlap
def generate_irregular_zones_no_overlap(rows, cols, num_zones, min_size, max_size):
    zone_polygons = []
    attempts = 0
    max_attempts = 100 * num_zones  # Prevent infinite loops

    while len(zone_polygons) < num_zones and attempts < max_attempts:
        center_x = np.random.randint(0, cols)
        center_y = np.random.randint(0, rows)
        size = np.random.randint(min_size, max_size)
        num_vertices = np.random.randint(3, 8)  # Polygons with 3 to 7 vertices

        # Generate random angles and radii for vertices
        angle = np.linspace(0, 2 * np.pi, num_vertices, endpoint=False)
        radius = size * np.random.rand(num_vertices)
        vertices = [
            (
                center_x + radius[i] * np.cos(angle[i]),
                center_y + radius[i] * np.sin(angle[i]),
            )
            for i in range(num_vertices)
        ]

        # Create the polygon
        polygon = Polygon(vertices)

        # Check for validity and overlap
        if polygon.is_valid and polygon.within(
            Polygon([(0, 0), (cols, 0), (cols, rows), (0, rows)])
        ):
            if not any(polygon.intersects(zone) for zone in zone_polygons):
                zone_polygons.append(polygon)

        attempts += 1

    if len(zone_polygons) < num_zones:
        print(
            f"Warning: Could only generate {len(zone_polygons)} zones due to overlap constraints."
        )

    return zone_polygons

# Function to generate the saturation map
def generate_saturation_map_from_zones_v3(
    rows, cols, zone_polygons, zone_saturations, retention_curve
):
    saturation_map = np.zeros((rows, cols))

    for polygon, saturation in zip(zone_polygons, zone_saturations):
        for r in range(rows):
            for c in range(cols):
                point = Point(c, r)
                if polygon.contains(point):
                    saturation_map[r, c] = saturation

    # Apply the water retention curve
    for j in range(cols):
        saturation_map[:, j] = apply_water_retention_curve_in_column_final(
            saturation_map[:, j], retention_curve
        )

    return saturation_map

# Function to calculate total water content
def calculate_total_water_content(saturation_map, cell_area):
    total_water_content = np.sum(saturation_map) * cell_area
    return total_water_content

# Archie's Law conversion function
def archies_law_conversion(saturation_map, phi, rho_w, theta_s, theta_r, a=1, m=2, n=2):
    # Ensure saturation values are within the valid range
    saturation_map = np.clip(saturation_map, theta_r, theta_s)
    # Convert volumetric water content to water saturation
    water_saturation = (saturation_map - theta_r) / (theta_s - theta_r)
    # Ensure saturation values are within the valid range for saturation
    water_saturation = np.clip(water_saturation, 0.001, 1)
    #set cover layer to be volumatrix water content to be 0.2, with 10% variation
    water_saturation[:,-1] = np.random.uniform(0,0.5)
    # Calculate resistivity using Archie's law
    resistivity_map = a * rho_w / (phi ** m * water_saturation ** n)

    return resistivity_map

def initialize_worker(seed):
    # Initialize random number generators with a unique seed
    np.random.seed(seed % (2**32 - 1))  # Ensure seed is within valid range
    # If we use the random module as well
    import random
    random.seed(seed % (2**32 - 1))




# %%
def perform_iteration(iteration):
    import numpy as np
    import pygimli as pg
    from pygimli.physics import ert
    # Optionally, reseed here as well to ensure uniqueness
    seed = (int(time.time() * 1000000) + os.getpid() + iteration) % (2**32 - 1)
    np.random.seed(seed)

    par_idx = np.random.randint(0,len(rr))
    # randomly sample a parameter set from the optimization result
    # Define all variables inside the function
    rows = 15
    cols = 142
    # Generate random number of zones, size, and saturation for each iteration
    # the range of free water from the model is around 0.25-1m, indicates around (0.25~1)/ (0.6288-0.386)= 1-4m in waste body is saturated.
    # 1/14 * 142 * 14 = 142 ~ 100
    # 4/14 * 142 *14 = 568 ~ 700
    # We choose these 100, 700 to allow some uncertainty
    num_zones = np.random.randint(100, 700)
    min_size = 1
    # Since the water balance model could not really detect the isolated water, we add some uncertainty to the upper limit 
    # Now the maximum saturated zones can be 3*700, slightly higher than total number of cells.
    # The overlapped parts will only be counted once in the generation procedure.
    max_size = 3
    zone_saturations = np.ones(num_zones)
    
    cell_area = 1.0  #  Area of each cell
    phi = 0.6288 + 0.2* np.random.randn() *  0.6288# Porosity 
    rho_w = 1.5 + 0.5 * np.random.randn() # Resistivity of pore water in ohm-meters
    rho_w = rho_w + (rho_w==0)*1.5
    a_arc = 1/ 10 **(optimized_par[par_idx,13])
    m_arc = optimized_par[par_idx,11]
    n_arc = optimized_par[par_idx,12]
    # Define the water retention curve for water materials (based on TDR experiment)
    theta_s = phi  # Saturated water content
    theta_r = optimized_par[par_idx,4]  # Residual water content
    alpha = optimized_par[par_idx,0]  # Inverse air entry suction
    n = optimized_par[par_idx,1] # Pore-size distribution index
    retention_curve = WaterRetentionCurve(theta_s=theta_s, theta_r=theta_r, alpha=alpha, n=n)

    # Define the grid for ERT simulation
    inversionDomain = pg.createGrid(x=np.arange(0, 143, 1), y=np.hstack((np.arange(-15.5, -0.5, 1),[0])), marker=2)
    grid = inversionDomain
    elecs = np.zeros((72, 2))
    elecs[:, 0] = np.arange(0, 144, 2)
    elecs[:, 1] = -0.2
    scheme_dd = ert.createData(elecs=elecs, schemeName='dd')
    spf.dd_generation(scheme_dd, max_geom_factor=5000)  
    scheme_slm = ert.createData(elecs=elecs, schemeName='slm')



    # Generate irregular zones and saturation maps
    zone_polygons = generate_irregular_zones_no_overlap(rows, cols, num_zones, min_size, max_size)
    # ! The saturation_map here is the Volumetric water content map, we transfer it to saturation map in function archies_law_conversion
    saturation_map = generate_saturation_map_from_zones_v3(rows, cols, zone_polygons, zone_saturations, retention_curve)

    # Calculate total water content
    total_water_content = calculate_total_water_content(saturation_map, cell_area)

    # Convert the saturation map to a resistivity map using Archie's law
    resistivity_map = archies_law_conversion(saturation_map, phi, rho_w, theta_s, theta_r, a=a_arc, m=m_arc, n=n_arc)

    # Simulate ERT data for both schemes (dd and slm)
    # Ensure that we extract only the numpy arrays to return
    data_dd_array = ert.simulate(grid, scheme=scheme_dd, res=resistivity_map.flatten(),
                                 noiseLevel=1, noiseAbs=1e-6, returnArray=True).array()

    data_slm_array = ert.simulate(grid, scheme=scheme_slm, res=resistivity_map.flatten(),
                                  noiseLevel=1, noiseAbs=1e-6, returnArray=True).array()

    # Return only pickleable objects (numpy arrays are pickleable)
    return (saturation_map, resistivity_map, total_water_content, data_dd_array, data_slm_array)

if __name__ == '__main__':
    from multiprocessing import Pool
    import os

    num_cores_available = os.cpu_count()
    print(f"Number of CPU cores available: {num_cores_available}")

    # Specify the number of cores to use
    num_cores_to_use = 20  # Adjust as needed

    # Number of iterations
    num_iterations = 100000

    # Use multiprocessing Pool to parallelize the loop
    seeds = [int(time.time() * 1000000) % 2**32 + i for i in range(num_cores_to_use)]
    with Pool(processes=num_cores_to_use, initializer=initialize_worker, initargs=(seeds[0],)) as pool:
        results = pool.map(perform_iteration, range(num_iterations))

    # Initialize lists to store the generated data
    saturation_maps = []
    resistivity_maps = []
    total_water_contents = []
    data_dd_list = []
    data_slm_list = []

    # Unpack the results
    for res in results:
        saturation_map, resistivity_map, total_water_content, data_dd_array, data_slm_array = res
        saturation_maps.append(saturation_map)
        resistivity_maps.append(resistivity_map)
        total_water_contents.append(total_water_content)
        data_dd_list.append(data_dd_array)
        data_slm_list.append(data_slm_array)

    # Convert lists to numpy arrays after the loop
    saturation_maps = np.array(saturation_maps)
    resistivity_maps = np.array(resistivity_maps)
    total_water_contents = np.array(total_water_contents)
    data_dd_array = np.array(data_dd_list)
    data_slm_array = np.array(data_slm_list)

    # Save all arrays to .npy files
    np.save("saturation_maps.npy", saturation_maps)
    np.save("resistivity_maps.npy", resistivity_maps)
    np.save("total_water_contents.npy", total_water_contents)
    np.save("data_dd.npy", data_dd_array)
    np.save("data_slm.npy", data_slm_array)