import numpy as np import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation from matplotlib.cm import ScalarMappable from matplotlib.colors import Normalize from IPython.display import Image # -------------------------------------------- # Constants and Initial Conditions # -------------------------------------------- g = 9.81 # m/s² Gamma_d = 9.8 / 1000 # Dry adiabatic lapse rate [K/m] z_break = 1000 # m, height of stable layer start Gamma_e1 = 6 / 1000 # 7 K/km in unstable layer Gamma_e2 = -5.0 / 1000 # -3 K/km in stable layer Gamma_dry = 9.8 / 1000 # Dry adiabatic lapse rate [K/m] T0 = 288 # Initial parcel temperature [K] Te = 285 # Initial environmental temperature [K] z0 = 1.0 # Initial altitude [m] (must be > z0_rough) w0 = 0.1 # Initial vertical velocity [m/s] r0 = 5.0 # Initial parcel radius [m] x0 = 30.0 # Initial x-position [m] # -------------------------------------------- # Log wind profile parameters # -------------------------------------------- u_star = 0.5 # friction velocity [m/s] kappa = 0.4 # von Kármán constant z0_rough = 0.1 # roughness length [m] # -------------------------------------------- # Time settings # -------------------------------------------- dt = 5.0 t_max = 1500 # seconds nt = int(t_max / dt) t = np.linspace(0, t_max, nt) # -------------------------------------------- #Buoyancy settings # -------------------------------------------- B_damping=0.01 # -------------------------------------------- # Initialize arrays # -------------------------------------------- z = np.zeros(nt) x = np.zeros(nt) T = np.zeros(nt) T_env = np.zeros(nt) w = np.zeros(nt) u = np.zeros(nt) N = np.zeros(nt) # sqrt of Brunt-Vaisala Ek = np.zeros(nt) # Energy Conservation radius = np.zeros(nt) # radius of the parcel size = np.zeros(nt) # size of the parcel dx = np.zeros(nt) # dynamic dx because of Langrangian approach CFL={} CFL['advection'] = np.zeros(nt) CFL['buoyancy'] = np.zeros(nt) # -------------------------------------------- # Initial values # -------------------------------------------- z[0] = z0 x[0] = x0 T[0] = T0 T_env[0] = Te w[0] = w0 radius[0] = r0 size[0] = (radius[0] * 4)**2 # -------------------------------------------- # Functions # -------------------------------------------- def EnergyConservation(w): return 1/2 * w**2 def BruntV(T, Gamma_e_current): N = np.sqrt(g/T*(Gamma_dry - Gamma_e_current)) return N def cfl(): # Estimate dx as parcel's horizontal travel in this step dx_local = x[n] - x[n-1] if n > 0 else 1.0 # avoid div by 0 dx[n] = dx_local # Compute CFL u_local = log_wind_profile(z[n]) CFL_local = u_local * dt / dx_local if dx_local != 0 else 0 return CFL_local def cfl_w(): dz_local = z[n] - z[n-1] if n > 0 else 1.0 # avoid div by 0 CFL_local = w[n] * dt / dz_local if dz_local != 0 else 0 return CFL_local def calc_equilibrium_level(T0, z0, Te, z_break, Gamma_d, Gamma_e1, Gamma_e2): # Temperature at break height for environmental profile T_break = Te - Gamma_e1 * z_break # EL candidate in first layer z_el_1 = (T0 - Te + Gamma_d * z0) / (Gamma_d - Gamma_e1) # EL candidate in second layer z_el_2 = (T0 - T_break + Gamma_d * z0 + Gamma_e2 * z_break) / (Gamma_d - Gamma_e2) # Check if EL in layer 1 is valid (between surface and z_break) if 0 <= z_el_1 <= z_break: return z_el_1 else: return z_el_2 def log_wind_profile(z, z0_rough=z0_rough, u_star=u_star, kappa=kappa): return u_star / kappa * np.log(z / z0_rough) def rk2_advection(x_old, z_old, dt): u1 = log_wind_profile(z_old) x_mid = x_old + 0.5 * u1 * dt z_mid = z_old u2 = log_wind_profile(z_mid) x_new = x_old + u2 * dt return x_new def rk4_step_w(w_prev, z_prev, T_prev, Te, z_break, Gamma_e1, Gamma_e2, dt, g=9.81, B_damping=0.0): def dwdt(w, z): # Compute environmental temperature if z < z_break: T_env = Te - Gamma_e1 * z else: T_break = Te - Gamma_e1 * z_break T_env = T_break - Gamma_e2 * (z - z_break) # Buoyancy with damping B = g * (T_prev - T_env) / T_env return B - B_damping * w k1 = dwdt(w_prev, z_prev) k2 = dwdt(w_prev + 0.5 * dt * k1, z_prev) k3 = dwdt(w_prev + 0.5 * dt * k2, z_prev) k4 = dwdt(w_prev + dt * k3, z_prev) return w_prev + (dt / 6.0) * (k1 + 2*k2 + 2*k3 + k4) el_height=calc_equilibrium_level(T0, z0, Te, z_break, Gamma_d, Gamma_e1, Gamma_e2) print(f"Equilibrium level height: {el_height:.2f} m") #Calculate T_env upfront since it is not advected z_array = np.maximum(z, 0) # ensure no negative heights T_break = Te - Gamma_e1 * z_break T_env = np.where( z_array < z_break, Te - Gamma_e1 * z_array, T_break - Gamma_e2 * (z_array - z_break) ) # -------------------------------------------- # Time Integration Loop # -------------------------------------------- for n in range(1, nt): # Select lapse rate depending on height Gamma_e_current = Gamma_e1 if z[n-1] < z_break else Gamma_e2 # Buoyancy B = g * (T[n-1] - T_env[n]) / T_env[n] if n == 1: # Euler step to initialize Leapfrog/RK4 w[n] = w[n-1] + B * dt z[n] = max(z[n-1] + w[n-1] * dt, z0) else: # Leapfrog scheme for vertical motion #w[n] = w[n-2] + 2 * B * dt #w[n] = w[n-2] + 2 * (B - B_damping * w[n-1]) * dt #z[n] = max(z[n-2] + 2 * w[n-1] * dt, z0) #RK4 for vertical motion w[n] = rk4_step_w(w[n-1], z[n-1], T[n-1], Te, z_break, Gamma_e1, Gamma_e2, dt, g, B_damping) z[n] = max(z[n-1] + w[n] * dt, z0) # Parcel temperature (Euler) T[n] = T[n-1] - Gamma_d * w[n-1] * dt # Horizontal position using RK2 x[n] = rk2_advection(x[n-1], z[n-1], dt) # Update radius and size radius[n] = r0 * (T0 / T[n])**(2) size[n] = (radius[n] * 4)**2 #acutal u u[n] = log_wind_profile(z[n]) CFL['advection'][n] = cfl() CFL['buoyancy'][n] = cfl_w() N[n]=BruntV(T[n], Gamma_e_current) # Normalize temperature for colormap norm = Normalize(vmin=min(T), vmax=max(T)) cmap = plt.cm.plasma ### Background stuff z_max = max(z) z_bg = np.linspace(0, z_max + 100, 300) # Compute environmental temperature profile T_break = Te - Gamma_e1 * z_break T_env_1D = np.where( z_bg < z_break, Te - Gamma_e1 * z_bg, T_break - Gamma_e2 * (z_bg - z_break) ) # Expand in x-direction x_bg = np.array([0, max(x) + 100]) X_bg, Z_bg = np.meshgrid(x_bg, z_bg) T_env_2D = np.tile(T_env_1D[:, np.newaxis], (1, len(x_bg))) ## Wind barbs # Define barb grid resolution z_levels = np.linspace(0, max(z), 20) x_positions = np.linspace(0, max(x), 5) Xb, Zb = np.meshgrid(x_positions, z_levels) # Get horizontal wind speed from profile (no vertical component) U_barb = np.array([log_wind_profile(zz)*1.1 for zz in z_levels]) U_barb = np.tile(U_barb[:, np.newaxis], (1, Xb.shape[1])) # repeat in x-direction W_barb = np.zeros_like(U_barb) # -------------------------------------------- # Animation Function # -------------------------------------------- def animate_growing_parcel(x, z, T, size, name='Parcel', parameters='Temp(euler)_RK2(wind)_RK4(buoyancy)'): filename = f'{name}_{parameters}.gif' nt = len(x) #dt = 1 # set this appropriately if you have a time step defined elsewhere fig, ax = plt.subplots(figsize=(8, 6)) ax.set_xlim(0, max(x) + 100) ax.set_ylim(0, max(z) + 100) ax.set_xlabel("Horizontal Position [m]") ax.set_ylabel("Altitude [m]") ax.set_title("Thermal Parcel Motion advected by Wind") # Temperature colormap normalization cmap = plt.cm.plasma norm = plt.Normalize(vmin=np.min(T), vmax=np.max(T)) trail = ax.scatter([], [], c=[], s=[], cmap=cmap, norm=norm) parcel_dot = ax.scatter([], [], c=[], s=[], edgecolors='k') sm = ScalarMappable(cmap=cmap, norm=norm) sm.set_array([]) cbar = plt.colorbar(sm, ax=ax) cbar.set_label("Temperature [K]") c_env = ax.contourf(X_bg, Z_bg, T_env_2D, levels=100, cmap=cmap, alpha=0.5, zorder=0) ax.barbs(Xb, Zb, U_barb, W_barb, length=5, barbcolor='gray', linewidth=0.7, alpha=0.9, zorder=1) # Equilibrium level line and label (assumes el_height is defined globally or passed in) el_height = np.max(z) * 0.75 # fallback if not defined el_line = ax.axhline(el_height, color='cyan', linestyle='--', label='Equilibrium Level') el_text = ax.text(0.02, el_height / ax.get_ylim()[1], "EL", transform=ax.transAxes, fontsize=12, color='cyan', va='bottom') # Add time display text time_text = ax.text(0.75, 0.95, '', transform=ax.transAxes, fontsize=12, bbox=dict(facecolor='white', alpha=0.7)) def init(): trail.set_offsets(np.empty((0, 2))) trail.set_array([]) trail.set_sizes([]) parcel_dot.set_offsets(np.empty((1, 2))) parcel_dot.set_sizes([0]) time_text.set_text('Time: 0 s') return trail, parcel_dot, time_text def update(i): coords = np.column_stack((x[:i+1], z[:i+1])) trail.set_offsets(coords) trail.set_array(T[:i+1]) trail.set_sizes(size[:i+1]) parcel_dot.set_offsets([[x[i], z[i]]]) parcel_dot.set_sizes([size[i]]) parcel_dot.set_array(np.array([T[i]])) time_text.set_text(f"Time: {i * dt } s") return trail, parcel_dot, time_text anim = FuncAnimation(fig, update, init_func=init, frames=nt, interval=100, blit=True) anim.save(filename, writer='pillow') ax.legend() plt.close(fig) return Image(filename=filename) # -------------------------------------------- # Run Animation # -------------------------------------------- animate_growing_parcel(x, z, T, size) #print(f"Max altitude reached: {max(z):.2f} m")

plt.plot(CFL['advection'][1:], label='Advection') plt.plot(CFL['buoyancy'][1:], label='Buoyancy') plt.xlabel('Steps') plt.ylabel('CFL') plt.title('CFL Stability check') plt.legend()

plt.plot(2/N, label='2/N') plt.axhline(dt, color='r', linestyle='--', label='dt') plt.xlabel('Steps') plt.ylabel('2 / Brunt-Vaisala Frequency') plt.title('Buoyancy Stability check') plt.legend()

def EnergyConservation(w,g,z,T): cp = 1005 return 1/2 * w**2 + g*z+ T*cp Ek = EnergyConservation(w,g,z,T) plt.plot(Ek/1000, label='Energy') plt.axhline(Ek[0]/1000, color='r', linestyle='--', label='Initial Energy') plt.legend() plt.xlabel('Steps') plt.ylabel('Energy [J/kg]') plt.title('Stability check by Energy Conservation')

plt.plot((Ek - Ek[0]) / Ek[0]) #plt.plot(Ek - Ek[0])

dt * u.max() / dx.max()

plt.plot(w)