from functools import partial
import numpy as np
from matplotlib import pyplot as plt
from scipy.integrate import solve_ivp
import scipy as sp
[docs]
class Viscosity:
def __init__(self, **kw):
for key in kw:
if not hasattr(key, self):
raise ValueError(f"Unkown {key=}")
setattr(self, kw[key])
self.init()
[docs]
def init(self):
self._V = sp.interpolate.InterpolatedUnivariateSpline(
[0, self.Lmid_L * self.L, self.L], [self.dV, self.V0, 0], k=1, ext="const"
)
self._dV = self._V.derivative()
[docs]
def E_V(self, n, d=0):
"""Return the internal (GPE) energy density and its derivatives."""
if d == 0:
return self.g * n**2 / 2
elif d == 1:
return self.g * n
elif d == 2:
return self.g + 0 * n
else:
return 0 * n
[docs]
def V(self, x, d=0):
"""Return the external potential or derivative."""
if d == 0:
return self._V(x)
elif d == 1:
return self._dV(x)
else:
raise NotImplementedError(f"Only d=1, 2 supported. Got {d=}.")
[docs]
def compute_dy_dt(self, x, y):
n, u, du, F = y
dn = -n * du / u
dV = self.V(x, d=1)
ddE = self.E_V(n, d=2)
ddu = n * (u * du + (dV + ddE * dn) / self.m) / self.nu
dF = n * dV
return (dn, du, ddu, dF)
[docs]
def solve(self, xL_L=-1, xR_L=1.2):
y0 = (self.n0, self.u0, 0, 0)
xL, xR = xL_L * self.L, xR_L * self.L
res = solve_ivp(self.compute_dy_dt, y0=y0, t_span=(xL, xR), max_step=0.001)
res.x = res.t
res.n, res.u, res.du, res.F = res.y
res.E, res.dE, res.ddE = [self.E_V(res.n, d=d) for d in [0, 1, 2]]
res.V, res.dV = [self.V(res.x, d=d) for d in [0, 1]]
res.P = res.n * (res.V + res.dE) - res.E
res.dn, res.du, res.ddu, res.F = self.compute_dy_dt(res.x, res.y)
assert np.allclose(res.n * res.u, self.n0 * self.u0)
res.p_current = res.n * res.u**2 + res.P / self.m - self.nu * res.du
res.n_current = res.n * res.u
return res
assert np.allclose(res.n_current[0], res.n_current)
assert np.allclose(res.p_current[0], res.p_current)
"""
fig, ax = plt.subplots()
ax.plot(x/L, n/n0, label="$n$")
ax.plot(x/L, u/u0, label="$u$")
ax.axvspan(0, 1, color='y', alpha=0.5)
ax.set(xlabel="$x/L$", ylabel="$n/n_0$", title=f"{n[0]=}, {n[-1]=}", );
plt.legend(loc="upper left")
ax = plt.twinx()
ax.plot(x/L, F, '--C1', label="$F$")
ax.legend(loc="lower right")
ax.set(ylabel="$F$");
"""