test.py

# %%
class WaterRetentionCurve:
    def __init__(self, theta_s, theta_r, alpha, n):
        """
        Initialize the Water Retention Curve using the van Genuchten model.

        Parameters:
            theta_s (float): Saturated water content.
            theta_r (float): Residual water content.
            alpha (float): Parameter related to the inverse of the air entry suction.
            n (float): Pore-size distribution index.
        """
        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):
        """
        Calculate water content given the pressure head h using the van Genuchten model.

        Parameters:
            h (float): Pressure head (positive downwards).

        Returns:
            theta (float): Water content at the given pressure head.
        """
        m = 1 - 1/self.n
        return self.theta_r + (self.theta_s - self.theta_r) / (1 + (self.alpha * abs(h))**self.n)**m
# %%
import numpy as np

def apply_water_retention_curve_in_column_final(column_data, retention_curve):
    """
    Apply the water retention curve to a column with multiple separated saturated zones,
    first detecting all saturated zones, and then handling unsaturated zones separately.

    Parameters:
        column_data (1D numpy array): A single column of the saturation map.
        retention_curve (WaterRetentionCurve): An instance of the WaterRetentionCurve class.

    Returns:
        column_water_content (1D numpy array): The updated column with water content applied only above each saturated zone.
    """
    column_water_content = np.full_like(column_data, retention_curve.theta_r)  # Start with residual water content

    # Step 1: Find all indices where the column is saturated (non-zero values)
    saturated_indices = np.nonzero(column_data)[0]
    column_water_content[saturated_indices] = 1
    # Step 2: Process each unsaturated zone separately
    if len(saturated_indices) > 0:
        # Unsaturated zone 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)

        # Unsaturated zones 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])

        # Unsaturated zone 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
# %%
import numpy as np
from shapely.geometry import Polygon, Point

def generate_irregular_zones_no_overlap(rows, cols, num_zones, min_size, max_size):
    """
    Generate non-overlapping irregular zones at arbitrary heights within a rectangular mesh.
    
    Parameters:
        rows (int): Number of rows in the mesh.
        cols (int): Number of columns in the mesh.
        num_zones (int): Number of zones to generate.
        min_size (int): Minimum size of the zones.
        max_size (int): Maximum size of the zones.
    
    Returns:
        zone_polygons (list of Polygon): List of Polygon objects representing the zones.
    """
    zone_polygons = []
    attempts = 0
    max_attempts = 100 * num_zones  # Limit the number of attempts to prevent infinite loops
    
    while len(zone_polygons) < num_zones and attempts < max_attempts:
        # Generate random center and size for the polygon
        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 if the polygon is within the bounds and doesn't overlap with existing polygons
        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
# %%
def generate_saturation_map_from_zones_v3(rows, cols, zone_polygons, zone_saturations, retention_curve):
    """
    Generate a 2D saturation map from irregular zones at arbitrary heights, applying the water retention curve to unsaturated areas.

    Parameters:
        rows (int): Number of rows in the mesh.
        cols (int): Number of columns in the mesh.
        zone_polygons (list of Polygon): List of Polygon objects representing the zones.
        zone_saturations (list of float): List of saturation levels corresponding to each zone.
        retention_curve (WaterRetentionCurve): The water retention curve for unsaturated zones.

    Returns:
        saturation_map (2D numpy array): The generated 2D saturation map.
    """
    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 new water retention curve logic
    for j in range(cols):
        saturation_map[:, j] = apply_water_retention_curve_in_column_final(saturation_map[:, j], retention_curve)

    return saturation_map
# %%
def archies_law_conversion(saturation_map, phi, rho_w, a=1, m=2, n=2):
    """
    Convert a saturation map to a resistivity map using Archie's law.

    Parameters:
        saturation_map (2D numpy array): The 2D map of water saturation values.
        phi (float): Porosity of the medium.
        rho_w (float): Resistivity of the pore water.
        a (float): Tortuosity factor (default is 1).
        m (float): Cementation exponent (default is 2).
        n (float): Saturation exponent (default is 2).

    Returns:
        resistivity_map (2D numpy array): The 2D resistivity map calculated from the saturation map.
    """
    resistivity_map = a * rho_w * (phi ** -m) * (saturation_map ** -n)
    return resistivity_map
# %%
import pygimli as pg
import pygimli.meshtools as mt
from pygimli.physics import ert

def simulate_ert_data(rows, cols, resistivity_map, electrode_spacing=1.0):
    """
    Simulate ERT data based on the resistivity map.
    
    Parameters:
        rows (int): Number of rows in the 2D map.
        cols (int): Number of columns in the 2D map.
        resistivity_map (2D numpy array): The resistivity map to simulate from.
        electrode_spacing (float): Spacing between electrodes.
    
    Returns:
        ert_data: Simulated ERT data.
    """
    # Create the world or geometry (adjust according to your specific case)
    world = mt.createWorld(start=[0, 0], end=[cols, -rows], layers=[-rows/2], worldMarker=True)

    # Generate the mesh
    mesh = mt.createMesh(world, quality=34, area=0.5)

    # Create electrode positions
    elecs = np.linspace(start=0, stop=cols, num=int(cols // electrode_spacing))

    # Create the ERT measurement scheme
    scheme = ert.createERTData(elecs=elecs, schemeName='dd')  # Dipole-dipole scheme

    # Simulate the ERT data
    ert_data = ert.simulate(mesh, res=resistivity_map, scheme=scheme, noiseLevel=1, noiseAbs=1e-6, seed=1337)
    
    return ert_data
# %%
def calculate_total_water_content(saturation_map, cell_area):
    """
    Calculate the total water content in the 2D saturation map.

    Parameters:
        saturation_map (2D numpy array): The 2D map of water saturation values.
        cell_area (float): The physical area represented by each cell.

    Returns:
        total_water_content (float): Total water content in the 2D map.
    """
    total_water_content = np.sum(saturation_map) * cell_area
    return total_water_content
# %%
import matplotlib.pyplot as plt

# Example settings
rows = 50
cols = 100
num_zones = 10
min_size = 5
max_size = 20
zone_saturations = np.ones(num_zones)

# Generate non-overlapping irregular zones with arbitrary heights
zone_polygons = generate_irregular_zones_no_overlap(rows, cols, num_zones, min_size, max_size)

# Define the water retention curve
theta_s = 1.0  # Saturated water content
theta_r = 0.1  # Residual water content
alpha = 0.03   # Inverse air entry suction
n = 2          # Pore-size distribution index
retention_curve = WaterRetentionCurve(theta_s=theta_s, theta_r=theta_r, alpha=alpha, n=n)

# Generate the saturation map using the updated logic
saturation_map = generate_saturation_map_from_zones_v3(rows, cols, zone_polygons, zone_saturations, retention_curve)

# Calculate the total water content
cell_area = 1.0  # Assume each cell represents an area of 1 square unit
total_water_content = calculate_total_water_content(saturation_map, cell_area)
print(f"Total Water Content: {total_water_content} cubic units")

# Optional: Visualize the saturation map
plt.imshow(saturation_map, cmap='Blues', origin='lower', extent=[0, cols, -rows, 0])
plt.colorbar(label='Saturation Level')
plt.title('Saturation Map with Corrected Irregular Zones')
plt.show()
# %%