<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
  
"""Shallow Water Model
  
- Two dimensional shallow water in a rotating frame
- Staggered Arakawa-C lat:lon grid
- periodic in the x-dimension
- fixed boundary conditions in the y-dimension
 
h = H + η
 
∂/∂t[u] - fv = - g ∂/∂x[η]
∂/∂t[v] + fu = - g ∂/∂y[η]
∂/∂t[h] + H(∂/∂x[u] + ∂/∂y[v]) = 0
 
f = f0 + βy
"""

import numpy as np

nx = 128
ny = 128

Lx = 100.0
Ly = 100.0

u = np.zeros((nx+3, ny+2))
v = np.zeros((nx+2, ny+3))
h = np.zeros((nx+2, ny+2))


f0 = 0.5
nu = 0.01
g = 9.8
H = 1.0

dx = dx = Lx / nx
dy = dy = Ly / ny

dt = 0.05
t = 0.0
tc = 0

def _update_boundaries(*vars):
    for var in vars:
    	# periodic x condition
        # copy values from lhs to the end of the array
        var[-2:, :] = var[1:3, :]
        # set the left hand boundary
        var[0, :] = var[-3, :]

        # zero deriv y condition
        var[:, 0] = var[:, 1]
        var[:, -1] = var[:, -2]

    	# fix corners to be average of neighbours
        var[0, 0] = 0.5*(var[1, 0] + var[0, 1])
        var[-1, 0] = 0.5*(var[-2, 0] + var[-1, 1])
        var[0, -1] = 0.5*(var[1, -1] + var[0, -2])
        var[-1, -1] = 0.5*(var[-1, -2] + var[-2, -1])


def diffx(psi):
    """Calculate ∂/∂x[psi] over a single grid square.
     
    i.e. d/dx(psi)[i,j] = (psi[i+1/2, j] - psi[i-1/2, j]) / dx
     
    The derivative is returned at x points at the midpoint between
    x points of the input array."""
    return (psi[1:,:] - psi[:-1,:]) / dx

def diff2x(psi):
    """Calculate ∂2/∂x2[psi] over a single grid square.
     
    i.e. d2/dx2(psi)[i,j] = (psi[i+1, j] - psi[i, j] + psi[i-1, j]) / dx^2
     
    The derivative is returned at the same x points as the
    x points of the input array, with dimension (nx-2, ny)."""
    return (psi[:-2, :] - psi[1:-1, :] + psi[2:, :]) / dx**2

def diff2y(psi):
    """Calculate ∂2/∂y2[psi] over a single grid square.
     
    i.e. d2/dy2(psi)[i,j] = (psi[i, j+1] - psi[i, j] + psi[i, j-1]) / dy^2
     
    The derivative is returned at the same y points as the
    y points of the input array, with dimension (nx, ny-2)."""
    return (psi[:, :-2] - psi[:, 1:-1] + psi[:, 2:]) / dy**2

def diffy(psi):
    """Calculate ∂/∂y[psi] over a single grid square.
     
    i.e. d/dy(psi)[i,j] = (psi[i, j+1/2] - psi[i, j-1/2]) / dy
     
    The derivative is returned at y points at the midpoint between
    y points of the input array."""
    return (psi[:, 1:] - psi[:,:-1]) / dy

def centre_average(phi):
    """Returns the four-point average at the centres between grid points."""
    return 0.25*(phi[:-1,:-1] + phi[:-1,1:] + phi[1:, :-1] + phi[1:,1:])
 
def y_average(phi):
    """Average adjacent values in the y dimension.
    If phi has shape (nx, ny), returns an array of shape (nx, ny - 1)."""
    return 0.5*(phi[:,:-1] + phi[:,1:])
 
def x_average(phi):
    """Average adjacent values in the x dimension.
    If phi has shape (nx, ny), returns an array of shape (nx - 1, ny)."""
    return 0.5*(phi[:-1,:] + phi[1:,:])

def divergence(u, v):
    """Returns the horizontal divergence at h points."""
    return diffx(u) + diffy(v)

def del2(phi):
    """Returns the Laplacian of phi."""
    return diff2x(phi)[:, 1:-1] + diff2y(phi)[1:-1, :]

def uvatuv(u, v):
    """Calculate the value of u at v and v at u."""
    ubar = centre_average(u)
    vbar = centre_average(v)
    return ubar, vbar

def uvath(u, v):
    ubar = x_average(u)
    vbar = y_average(v)
    return ubar, vbar


def rhs(state):
    """Calculate the right hand side of the u, v and h equations."""
    global f0, g, H
    u, v, h = state[0:3]
    uu, vv = uvatuv(u, v)
    
    u_rhs = np.zeros_like(u)
    v_rhs = np.zeros_like(v)
    h_rhs = np.zeros_like(h)

    # the height equation
    h_rhs[:] = -H*(divergence(u, v))

    # the u equation
    dhdx = diffx(h)
    u_rhs[1:-1, :] = f0*vv - g*dhdx 
    u_rhs[1:-1, 1:-1] += nu*del2(u)
     
    # the v equation
    dhdy = diffy(h)
    v_rhs[:, 1:-1] = -f0*uu - g*dhdy
    v_rhs[1:-1, 1:-1] += nu*del2(v)

    return np.array([u_rhs, v_rhs, h_rhs])




_ppdstate, _pdstate = 0,0
def step(state):
    global dt, t, tc, _ppdstate, _pdstate

    _update_boundaries(*state)
    
    dstate = rhs(state)

    # take adams-bashforth step in time
    if tc==0:
        # forward euler
        dt1 = dt
        dt2 = 0.0
        dt3 = 0.0
    elif tc==1:
        # AB2 at step 2
        dt1 = 1.5*dt
        dt2 = -0.5*dt
        dt3 = 0.0
    else:
        # AB3 from step 3 on
        dt1 = 23./12.*dt
        dt2 = -16./12.*dt
        dt3 = 5./12.*dt
    
    newstate = state + dt1*dstate + dt2*_pdstate + dt3*_ppdstate
    state = newstate
    _ppdstate = _pdstate
    _pdstate = dstate

    t  += dt
    tc += 1
    return state

import matplotlib.pyplot as plt

plt.figure(figsize=(12, 5))
def plot_all(u,v,h):
    plt.clf()
    plt.subplot(131)
    plt.imshow(u.T)
    plt.subplot(132)
    plt.imshow(v.T)
    plt.subplot(133)
    plt.imshow(h.T, cmap=plt.cm.seismic)
    plt.clim(-np.abs(h).max(), np.abs(h).max())
    #plt.colorbar(orientation='horizontal')

    plt.pause(0.1)


h[70:80, 70:80] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
h[40:50, 40:50] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
h[100:110, 70:80] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
#h[48:128, 48:128] = np.sin(np.linspace(0, np.pi, 80))[:, np.newaxis]**2*np.sin(np.linspace(0, np.pi, 80))[np.newaxis, :]**2*0.01
#h[np.sum((np.indices(h.shape)- 60)**2) < 300] = 0.01
state = np.array([u, v, h])


for i in range(10000):
    state = step(state)
    if i % 20 == 0:
        plot_all(*state)
        print(state[2].min(), state[2].max())<commit_msg>Added beta term and fixed node positions<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
  
"""Shallow Water Model
  
- Two dimensional shallow water in a rotating frame
- Staggered Arakawa-C lat:lon grid
- periodic in the x-dimension
- fixed boundary conditions in the y-dimension
 
h = H + η
 
∂/∂t[u] - fv = - g ∂/∂x[η]
∂/∂t[v] + fu = - g ∂/∂y[η]
∂/∂t[h] + H(∂/∂x[u] + ∂/∂y[v]) = 0
 
f = f0 + βy
"""

import numpy as np

nx = 128
ny = 129

Lx = 1.0e5
Ly = 1.0e5

u = np.zeros((nx+3, ny+2))
v = np.zeros((nx+2, ny+3))
h = np.zeros((nx+2, ny+2))

dx = Lx / nx
dy = Ly / ny

# positions of the nodes
ux = (-Lx/2 + np.arange(nx+1)*dx)[:, np.newaxis]
vx = (-Lx/2 + dx/2.0 + np.arange(nx)*dx)[:, np.newaxis]

vy = (-Ly/2 + np.arange(ny+1)*dy)[np.newaxis, :]
uy = (-Ly/2 + dy/2.0 + np.arange(ny)*dy)[np.newaxis, :]

hx = vx
hy = uy


f0 = 0.0
beta = 1e-6
nu = 0.1   # diffusion
g = 9.8
H = 10.0

dt = 2.0
t = 0.0
tc = 0


def _update_boundaries(*vars):
    for var in vars:
    	# periodic x condition
        # copy values from lhs to the end of the array
        var[-2:, :] = var[1:3, :]
        # set the left hand boundary
        var[0, :] = var[-3, :]

        # zero deriv y condition
        var[:, 0] = var[:, 1]
        var[:, -1] = var[:, -2]

    	# fix corners to be average of neighbours
        var[0, 0] = 0.5*(var[1, 0] + var[0, 1])
        var[-1, 0] = 0.5*(var[-2, 0] + var[-1, 1])
        var[0, -1] = 0.5*(var[1, -1] + var[0, -2])
        var[-1, -1] = 0.5*(var[-1, -2] + var[-2, -1])


def diffx(psi):
    """Calculate ∂/∂x[psi] over a single grid square.
     
    i.e. d/dx(psi)[i,j] = (psi[i+1/2, j] - psi[i-1/2, j]) / dx
     
    The derivative is returned at x points at the midpoint between
    x points of the input array."""
    return (psi[1:,:] - psi[:-1,:]) / dx

def diff2x(psi):
    """Calculate ∂2/∂x2[psi] over a single grid square.
     
    i.e. d2/dx2(psi)[i,j] = (psi[i+1, j] - psi[i, j] + psi[i-1, j]) / dx^2
     
    The derivative is returned at the same x points as the
    x points of the input array, with dimension (nx-2, ny)."""
    return (psi[:-2, :] - psi[1:-1, :] + psi[2:, :]) / dx**2

def diff2y(psi):
    """Calculate ∂2/∂y2[psi] over a single grid square.
     
    i.e. d2/dy2(psi)[i,j] = (psi[i, j+1] - psi[i, j] + psi[i, j-1]) / dy^2
     
    The derivative is returned at the same y points as the
    y points of the input array, with dimension (nx, ny-2)."""
    return (psi[:, :-2] - psi[:, 1:-1] + psi[:, 2:]) / dy**2

def diffy(psi):
    """Calculate ∂/∂y[psi] over a single grid square.
     
    i.e. d/dy(psi)[i,j] = (psi[i, j+1/2] - psi[i, j-1/2]) / dy
     
    The derivative is returned at y points at the midpoint between
    y points of the input array."""
    return (psi[:, 1:] - psi[:,:-1]) / dy

def centre_average(phi):
    """Returns the four-point average at the centres between grid points."""
    return 0.25*(phi[:-1,:-1] + phi[:-1,1:] + phi[1:, :-1] + phi[1:,1:])
 
def y_average(phi):
    """Average adjacent values in the y dimension.
    If phi has shape (nx, ny), returns an array of shape (nx, ny - 1)."""
    return 0.5*(phi[:,:-1] + phi[:,1:])
 
def x_average(phi):
    """Average adjacent values in the x dimension.
    If phi has shape (nx, ny), returns an array of shape (nx - 1, ny)."""
    return 0.5*(phi[:-1,:] + phi[1:,:])

def divergence(u, v):
    """Returns the horizontal divergence at h points."""
    return diffx(u) + diffy(v)

def del2(phi):
    """Returns the Laplacian of phi."""
    return diff2x(phi)[:, 1:-1] + diff2y(phi)[1:-1, :]

def uvatuv(u, v):
    """Calculate the value of u at v and v at u."""
    ubar = centre_average(u)[1:-1, :]
    vbar = centre_average(v)[:, 1:-1]
    return ubar, vbar

def uvath(u, v):
    ubar = x_average(u)
    vbar = y_average(v)
    return ubar, vbar


def rhs(state):
    """Calculate the right hand side of the u, v and h equations."""
    global f0, g, H
    u, v, h = state[0:3]
    uu, vv = uvatuv(u, v)
    
    u_rhs = np.zeros_like(u)
    v_rhs = np.zeros_like(v)
    h_rhs = np.zeros_like(h)

    # the height equation
    h_rhs[:] = -H*(divergence(u, v))

    # the u equation
    dhdx = diffx(h)[:, 1:-1]
    u_rhs[1:-1, 1:-1] = f0*vv + beta*uy*vv - g*dhdx + nu*del2(u)*np.abs(uy/Ly)
     
    # the v equation
    dhdy = diffy(h)[1:-1, :]
    v_rhs[1:-1, 1:-1] = -f0*uu - beta*vy*uu - g*dhdy + nu*del2(v)*np.abs(vy/Ly)

    return np.array([u_rhs, v_rhs, h_rhs])




_ppdstate, _pdstate = 0,0
def step(state):
    global dt, t, tc, _ppdstate, _pdstate

    _update_boundaries(*state)
    
    dstate = rhs(state)

    # take adams-bashforth step in time
    if tc==0:
        # forward euler
        dt1 = dt
        dt2 = 0.0
        dt3 = 0.0
    elif tc==1:
        # AB2 at step 2
        dt1 = 1.5*dt
        dt2 = -0.5*dt
        dt3 = 0.0
    else:
        # AB3 from step 3 on
        dt1 = 23./12.*dt
        dt2 = -16./12.*dt
        dt3 = 5./12.*dt
    
    newstate = state + dt1*dstate + dt2*_pdstate + dt3*_ppdstate
    state = newstate
    _ppdstate = _pdstate
    _pdstate = dstate

    t  += dt
    tc += 1
    return state

import matplotlib.pyplot as plt

plt.figure(figsize=(12, 5))
def plot_all(u,v,h):
    plt.clf()
    plt.subplot(131)
    plt.imshow(u[1:-1, 1:-1].T,
            extent=[ux.min(), ux.max(), uy.min(), uy.max()])
    plt.subplot(132)
    plt.imshow(v[1:-1, 1:-1].T,
            extent=[vx.min(), vx.max(), vy.min(), vy.max()])
    plt.subplot(133)
    plt.imshow(h[1:-1, 1:-1].T, cmap=plt.cm.seismic,
            extent=[hx.min(), hx.max(), hy.min(), hy.max()])
    plt.clim(-np.abs(h).max(), np.abs(h).max())
    #plt.colorbar(orientation='horizontal')

    plt.pause(0.1)


h[nx//2-5:nx//2+5, ny//2-5:ny//2+5] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
#h[70:80, 70:80] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
#h[40:50, 40:50] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
#h[100:110, 70:80] = np.exp(-(((np.indices((10,10)) - 5)/2.0)**2).sum(axis=0))
#h[48:128, 48:128] = np.sin(np.linspace(0, np.pi, 80))[:, np.newaxis]**2*np.sin(np.linspace(0, np.pi, 80))[np.newaxis, :]**2*0.01
#h[np.sum((np.indices(h.shape)- 60)**2) < 300] = 0.01
state = np.array([u, v, h])


for i in range(100000):
    state = step(state)
    if i % 20 == 0:
        plot_all(*state)
        print(state[2].min(), state[2].max())
{'fraction_non_alphanumeric': 0.20505175983436852, 'fraction_numerical': 0.055900621118012424, 'mean_word_length': 2.369419642857143, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5082590', 'n_tokens_mistral': 5543, 'n_tokens_neox': 5244, 'n_words': 1348}