pip install landlab matplotlib numpy

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from landlab import RasterModelGrid
from landlab.components import FlowAccumulator, FastscapeEroder, LinearDiffuser
from landlab.plot import imshow_grid

# 1. Setup Grid
mg = RasterModelGrid((50, 50), xy_spacing=10.0)
z = mg.add_zeros("topographic__elevation", at="node")
# Add some initial random noise to the landscape
z += np.random.rand(mg.number_of_nodes) * 0.1

# 2. Boundary Conditions
# Set edges to closed, except for the bottom (outlet)
mg.set_closed_boundaries_at_grid_edges(True, True, True, False)

# 3. Initialize Components
fr = FlowAccumulator(mg, flow_director='FlowDirectorD8')
sp = FastscapeEroder(mg, K_sp=0.0001, m_sp=0.5, n_sp=1.0)
ld = LinearDiffuser(mg, linear_diffusivity=0.01)

# 4. Simulation Parameters
dt = 1000
total_time = 500000
uplift_rate = 0.0001
n_steps = int(total_time // dt)

# 5. Animation Setup
fig, ax = plt.subplots()
artists = []

# 6. Run Model
for i in range(n_steps):
    # Uplift
    z[mg.core_nodes] += uplift_rate * dt
    
    # Process
    fr.run_one_step()
    sp.run_one_step(dt)
    ld.run_one_step(dt)
    
    # Store frames
    if i % 20 == 0:
        im = imshow_grid(mg, z, cmap='terrain', allow_colorbar=False, plot_name=f"Time: {i*dt}")
        artists.append([im])

# 7. Generate Animation
ani = animation.ArtistAnimation(fig, artists, interval=50, blit=True)
ani.save('landscape_evolution.gif', writer='pillow')
plt.show()

pip install landlab numpy matplotlib

import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
from landlab import RasterModelGrid
from landlab.components import FlowAccumulator, LinearDiffuser, StreamPowerEroder

# ==========================================
# 1. Setup Grid and Parameters
# ==========================================
num_rows = 50
num_cols = 50
dx = 10.0  # Grid spacing in meters

grid = RasterModelGrid((num_rows, num_cols), xy_spacing=dx)

# Create the elevation field and add random noise to initiate drainage
z = grid.add_zeros("topographic__elevation", at="node")
np.random.seed(42)
z += np.random.rand(grid.number_of_nodes) * 0.1

# Set boundary conditions: Bottom edge open, others closed
grid.set_closed_boundaries_at_grid_edges(
    right_is_closed=True,
    top_is_closed=True,
    left_is_closed=True,
    bottom_is_closed=False,
)

# Time parameters
dt = 100.0  # Timestep in years
total_steps = 200  # Number of frames in animation
uplift_rate = 0.001  # Uplift rate in m/yr

# Component parameters
K_sp = 0.001  # Stream power erodibility coefficient
m_sp = 0.5  # Area exponent
n_sp = 1.0  # Slope exponent
D = 0.01  # Hillslope diffusivity (m^2/yr)

# ==========================================
# 2. Initialize Landlab Components
# ==========================================
fa = FlowAccumulator(grid, surface="topographic__elevation", flow_director="D8")
sp = StreamPowerEroder(grid, K_sp=K_sp, m_sp=m_sp, n_sp=n_sp)
ld = LinearDiffuser(grid, linear_diffusivity=D)

# ==========================================
# 3. Setup Animation Plotting
# ==========================================
fig, ax = plt.subplots(figsize=(8, 6))

# Reshape elevation for 2D plotting
z_matrix = z.reshape((num_rows, num_cols))

# Initial plot setup
im = ax.imshow(
    z_matrix,
    cmap="terrain",
    extent=[0, num_cols * dx, 0, num_rows * dx],
    origin="lower",
)
fig.colorbar(im, ax=ax, label="Elevation (m)")
title = ax.set_title("Landscape Evolution: Year 0")
ax.set_xlabel("X Distance (m)")
ax.set_ylabel("Y Distance (m)")


# ==========================================
# 4. Animation Update Function
# ==========================================
def update_landscape(frame):
    # 1. Apply tectonic uplift to core (interior) nodes
    z[grid.core_nodes] += uplift_rate * dt

    # 2. Route flow across the landscape
    fa.run_one_step()

    # 3. Fluvial incision (Stream Power)
    sp.run_one_step(dt)

    # 4. Hillslope creep (Diffusion)
    ld.run_one_step(dt)

    # Update visual data
    z_matrix = z.reshape((num_rows, num_cols))
    im.set_data(z_matrix)
    im.set_clim(vmin=z_matrix.min(), vmax=z_matrix.max())
    title.set_text(f"Landscape Evolution: Year {int((frame + 1) * dt)}")

    return [im, title]


# ==========================================
# 5. Run and Save Animation
# ==========================================
ani = animation.FuncAnimation(
    fig,
    update_landscape,
    frames=total_steps,
    interval=50,  # Delay between frames in ms
    blit=False,
)

# To view live while running, uncomment the line below:
# plt.show()

# Save the animation as an MP4 file (requires ffmpeg)
output_filename = "landscape_evolution.mp4"
print(f"Running simulation and saving to {output_filename}...")
ani.save(output_filename, writer="ffmpeg", fps=20)
print("Done!")

#!/usr/bin/env python3
"""
Stream Power + Diffusion (SPD) landscape evolution model in LandLab.
Runs a 50x50 grid, evolves the topography, and generates an animation.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.colors import LightSource
from landlab import RasterModelGrid
from landlab.components import FlowRouter, StreamPowerEroder, LinearDiffuser
from landlab.io.netcdf import write_netcdf


def run_spd_model(
    grid_size=50,
    n_steps=500,
    dt=100,
    k_sp=1e-4,  # Stream power erodibility
    k_diff=0.01,  # Diffusivity coefficient
    m=0.5,  # Stream power slope exponent
    n=1.0,  # Stream power discharge exponent
):
    """
    Run SPD landscape evolution model.
    
    Parameters
    ----------
    grid_size : int
        Grid dimensions (grid_size x grid_size)
    n_steps : int
        Number of model iterations
    dt : float
        Time step in years
    k_sp : float
        Stream power erodibility coefficient
    k_diff : float
        Hillslope diffusivity
    m : float
        Stream power slope exponent
    n : float
        Stream power discharge exponent
    
    Returns
    -------
    grid : RasterModelGrid
        Final model grid
    elevation_history : list
        Elevation arrays at each saved timestep
    time_steps : list
        Corresponding model times (years)
    """
    
    # Create grid
    grid = RasterModelGrid((grid_size, grid_size), xy_spacing=50.0)
    
    # Add elevation field (initial condition: gentle noise on a slope)
    np.random.seed(42)
    x, y = grid.x_of_node, grid.y_of_node
    base_elev = 1000 + 0.001 * y  # Gentle slope towards north
    noise = np.random.normal(0, 5, grid.number_of_nodes)
    elevation = base_elev + noise
    grid.add_field("topographic__elevation", elevation, at="node")
    
    # Set boundary conditions (fixed elevation at edges)
    grid.set_watershed_boundary_condition_outlet_id(
        0, grid.nodes[0, 0], -9999
    )
    
    # Initialize components
    flow_router = FlowRouter(grid)
    sp_eroder = StreamPowerEroder(grid, K=k_sp, m=m, n=n, threshold_sp=0)
    diffuser = LinearDiffuser(grid, linear_diffusivity=k_diff)
    
    # Storage for animation
    elevation_history = []
    time_steps = []
    save_interval = max(1, n_steps // 50)  # Save ~50 frames
    
    current_time = 0
    
    print(f"Running {n_steps} steps on a {grid_size}x{grid_size} grid...")
    
    for step in range(n_steps):
        # Route flow
        flow_router.run_one_step()
        
        # Erode with stream power
        sp_eroder.run_one_step(dt)
        
        # Diffuse hillslopes
        diffuser.run_one_step(dt)
        
        current_time += dt
        
        # Save frame for animation
        if step % save_interval == 0:
            elevation_history.append(grid["topographic__elevation"].copy())
            time_steps.append(current_time)
        
        if (step + 1) % 50 == 0:
            print(f"  Step {step + 1}/{n_steps} (t={current_time/1000:.1f} ky)")
    
    print(f"Model complete. Final time: {current_time/1000:.1f} ky")
    return grid, elevation_history, time_steps


def create_animation(elevation_history, time_steps, output_file="spd_evolution.mp4"):
    """
    Create and save animation of landscape evolution.
    
    Parameters
    ----------
    elevation_history : list
        List of elevation arrays
    time_steps : list
        List of corresponding times
    output_file : str
        Output video filename
    """
    
    fig, ax = plt.subplots(figsize=(10, 9))
    
    # Set up color limits for consistent scaling
    vmin = min(arr.min() for arr in elevation_history)
    vmax = max(arr.max() for arr in elevation_history)
    
    # Reshape for 2D visualization
    grid_size = int(np.sqrt(len(elevation_history[0])))
    
    im = None
    
    def update(frame):
        nonlocal im
        
        # Reshape elevation to 2D grid
        elev_2d = elevation_history[frame].reshape((grid_size, grid_size))
        
        # Clear previous image
        if im is not None:
            im.remove()
        
        # Plot with hillshade for better visualization
        ls = LightSource(azdeg=315, altdeg=45)
        hillshade = ls.hillshade(elev_2d, vert_exag=2)
        
        im = ax.imshow(
            hillshade,
            cmap="gray",
            origin="lower",
            extent=[0, grid_size * 50, 0, grid_size * 50],
        )
        
        # Add elevation contours
        contours = ax.contour(
            elev_2d,
            levels=10,
            colors="white",
            alpha=0.3,
            extent=[0, grid_size * 50, 0, grid_size * 50],
            linewidths=0.5,
        )
        
        time_ka = time_steps[frame] / 1000
        ax.set_title(
            f"SPD Landscape Evolution\nTime: {time_ka:.1f} ka",
            fontsize=14,
            fontweight="bold",
        )
        ax.set_xlabel("Distance (m)")
        ax.set_ylabel("Distance (m)")
        
        return [im]
    
    anim = animation.FuncAnimation(
        fig,
        update,
        frames=len(elevation_history),
        interval=50,
        blit=True,
        repeat=True,
    )
    
    print(f"Saving animation to {output_file}...")
    anim.save(output_file, writer="ffmpeg", dpi=100)
    print(f"Animation saved!")
    
    plt.close(fig)


def plot_summary(grid, elevation_history):
    """Plot initial and final topography."""
    
    grid_size = int(np.sqrt(len(elevation_history[0])))
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    elev_init = elevation_history[0].reshape((grid_size, grid_size))
    elev_final = elevation_history[-1].reshape((grid_size, grid_size))
    
    im0 = axes[0].imshow(elev_init, cmap="viridis", origin="lower")
    axes[0].set_title("Initial Topography", fontweight="bold")
    axes[0].set_xlabel("Grid index")
    axes[0].set_ylabel("Grid index")
    plt.colorbar(im0, ax=axes[0], label="Elevation (m)")
    
    im1 = axes[1].imshow(elev_final, cmap="viridis", origin="lower")
    axes[1].set_title("Final Topography", fontweight="bold")
    axes[1].set_xlabel("Grid index")
    axes[1].set_ylabel("Grid index")
    plt.colorbar(im1, ax=axes[1], label="Elevation (m)")
    
    # Add relief and slope info
    relief_init = elev_init.max() - elev_init.min()
    relief_final = elev_final.max() - elev_final.min()
    
    fig.text(
        0.5,
        0.02,
        f"Initial relief: {relief_init:.1f} m | Final relief: {relief_final:.1f} m",
        ha="center",
        fontsize=11,
    )
    
    plt.tight_layout()
    plt.subplots_adjust(bottom=0.08)
    plt.savefig("spd_summary.png", dpi=150, bbox_inches="tight")
    print("Summary figure saved to spd_summary.png")
    plt.close(fig)


if __name__ == "__main__":
    # Run model
    grid, elev_hist, times = run_spd_model(
        grid_size=50,
        n_steps=500,
        dt=100,
        k_sp=1e-4,
        k_diff=0.01,
        m=0.5,
        n=1.0,
    )
    
    # Create animation
    create_animation(elev_hist, times, "spd_evolution.mp4")
    
    # Plot summary
    plot_summary(grid, elev_hist)
    
    print("\nDone! Check spd_evolution.mp4 and spd_summary.png")

python spd_animation.py

pip install landlab matplotlib numpy ffmpeg-python