<jupyter_start><jupyter_text>Two-layer code comparison problem<jupyter_code>%pylab widget
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp, cumtrapz

import logging
logger = logging.getLogger()
logger.setLevel(logging.WARNING)

# Some of Paul Tol's colourblind-safe colours.
lc = ('#0077bb', '#009988', '#ee7733', '#cc3311')
Rgasconst = 1.108800

# Reference point.
y0 = 1.

# Density and pressure at the reference point.
rho0 = 1.
p0 = 0.6

# Mean molecular weights of the two fluids.
mu0 = 1.848
mu1 = 1.802

# gamma = dln(p)/dln(rho) for the two layers.
gamma0 = 5./3.
gamma1 = 1.3
def fg_func(y):
    """
    Returns the function fg(y) that turns of gravity at the
    top and bottom boundaries.
    """
    
    if y < 1.:
        return 0.
    elif y < 1.0625:
        return 0.5*(1. + np.sin(16.*np.pi*(y - 1.03125)))
    elif y < 2.9375:
        return 1.
    elif y < 3.:
        return 0.5*(1. - np.sin(16.*np.pi*(y - 2.96875)))
    else:
        return 0.

def g_func(y):
    """
    Returns the gravity profile g(y).
    """
    
    g0 = -1.414870
    return fg_func(y)*g0*y**(-5./4.)

def fv_func(y):
    """
    Returns the profile fv(y) of the fractional volume of
    the lighter fluid.
    """
    
    if y < 1.9375:
        return 0.
    elif y < 2.0625:
        return 0.5*(1. + np.sin(8.*np.pi*(y - 2.)))
    else:
        return 1.

def dUdy(y, U):
    """
    Returns the right hand sides of the two differential
    equations that define the stratification:

    dln(rho)/dy = (dln(p)/dy)/gamma,
    dln(p)/dy = rho*g/p,
    
    where gamma = gamma0 + fv_func(y)*(gamma1 - gamma0) and
    we are integrating for ln(rho) and ln(p).
    """

    dlnpdy = np.exp(U[0])*g_func(y)/np.exp(U[1])
    gamma = gamma0 + fv_func(y)*(gamma1 - gamma0)
    dlnrhody = dlnpdy/gamma
    
    return [dlnrhody, dlnpdy]

def cdiff(x):
    """
    Returns the 2nd order central difference with one-sided
    1st order differences applied at the boundaries of the
    array x.
    """

    dx = 0.5*(np.roll(x, -1) - np.roll(x, +1))
    dx[0] = x[1] - x[0]
    dx[-1] = x[-1] - x[-2]

    return dx
# Vertical coordinate range. The computational domain
# is (1., 3.) and the rest is provided to fill ghost
# cells if needed.
ylim = (0.9, 3.1)
ngp = 10001
y = np.linspace(ylim[0], ylim[1], ngp)
dy = y[1] - y[0]

# Gravity.
g = np.vectorize(g_func)(y)

# Gravitational potential.
gpot = -cumtrapz(g, x=y, initial=0.)
idx0 = np.argmin(np.abs(y - 3.))
gpot -= gpot[idx0]

# Fractional volume of the stably-stratified fluid.
fv = np.vectorize(fv_func)(y)

# Mean molecular weight.
mu = (1. - fv)*mu0 + fv*mu1

idx0 = np.argmax(y > y0)

# Initial condition.
U0 = [np.log(rho0), np.log(p0)]

# Solve for HSE.
U_above = solve_ivp(dUdy, (y0, ylim[1]), U0, t_eval=y[idx0:], \
                    max_step=dy, method='RK45', atol=1e-6)
U_below = solve_ivp(dUdy, (y0, ylim[0]), U0, t_eval=np.flip(y[:idx0]), \
                    max_step=dy, method='RK45', atol=1e-6)

lnrho = np.concatenate((np.flip(U_below['y'][0]), U_above['y'][0]))
lnp = np.concatenate((np.flip(U_below['y'][1]), U_above['y'][1]))
lnT = np.log(np.exp(lnp)*mu/(Rgasconst*np.exp(lnrho)))<jupyter_output><empty_output><jupyter_text>Input gravity and fractional volume of the stably-stratified fluid<jupyter_code>ifig=1; plt.close(ifig); plt.figure()

ax0 = plt.gca()
ax0.plot(y, g, '-', color=lc[1], label=r'$g$')
ax0.legend(loc=2)
ax0.set_xlim(ylim)
ax0.set_ylim((-1.4, 0.1))
ax0.set_xlabel(r'$y$')
ax0.set_ylabel(r'$g$')

ax1 = ax0.twinx()
ax1.plot(y, fv, '--', color=lc[2], label=r'fv')
ax1.legend(loc=4)
ax1.set_xlim(ylim)
ax1.set_ylim((-0.05, 1.05))
ax1.set_xlabel(r'$y$')
ax1.set_ylabel(r'fv')<jupyter_output><empty_output><jupyter_text>Hydrostatic stratification of density, pressure, temperature, and mean molecular weight<jupyter_code>ifig=2; plt.close(ifig); plt.figure()

ax0 = plt.gca()
ax0.plot(y, lnrho, '-', color=lc[0], label=r'$\ln(\rho)$')
ax0.plot(y, lnp, '--', color=lc[1], label=r'$\ln(p)$')
ax0.plot(y, lnT, '-.', color=lc[2], label=r'$\ln(T)$')
ax0.legend(loc=3)
ax0.set_xlim(ylim)
ax0.set_ylim((-5.25, 0.25))
ax0.set_xlabel(r'$y$')
ax0.set_ylabel(r'$\ln(\rho)$, $\ln(p)$, $\ln(T)$')

ax1 = ax0.twinx()
ax1.plot(y, mu, ':', color=lc[3], label=r'$\mu$')
ax1.legend(loc=1)
ax1.set_xlim(ylim)
ax1.set_ylim((1.80, 1.85))
ax1.set_xlabel(r'$y$')
ax1.set_ylabel(r'$\mu$')<jupyter_output><empty_output><jupyter_text>Check that the stratification is hydrostatic<jupyter_code>ifig=3; plt.close(ifig); plt.figure()
hse_error = cdiff(lnp)/cdiff(y) - np.exp(lnrho)*g/np.exp(lnp)
plt.semilogy(y, np.abs(hse_error), '-', color=lc[0])
plt.xlim(ylim)
plt.xlabel(r'$y$')
plt.ylabel(r'$|\mathrm{d}\ln(p) / \mathrm{d}y - \rho g / p|$')<jupyter_output><empty_output><jupyter_text>Check the pressure-density relation<jupyter_code>ifig=4; plt.close(ifig); plt.figure()
gamma = cdiff(lnp)/cdiff(lnrho)
plt.plot(y, gamma, '-', color=lc[0])
plt.xlim(ylim)
plt.xlabel(r'$y$')
plt.ylabel(r'$\mathrm{d}\ln(p) / \mathrm{d}\ln(\rho)$')<jupyter_output><empty_output><jupyter_text>Write the stratification into a file<jupyter_code>with open('setup-two-layers.in', 'w') as fout:
    fout.write('{:d}\n'.format(ngp))
    fstr = '{:.6e}   '*6 + '\n'
    for row in zip(y, g, gpot, fv, np.exp(lnrho), np.exp(lnp)):
        fout.write(fstr.format(*row))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19597989949748743, 'fraction_numerical': 0.04690117252931323, 'mean_word_length': 3.612875536480687, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 26, '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}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '416319', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2492, 'n_tokens_neox': 2362, 'n_words': 469, 'path': '1D2D/setup-two-layers.ipynb', 'repo_name': 'robert-andrassy/CoCoPy', 'size': '5373'}