## Stommel/ Taylor model
## Eli Tziperman, 201910
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from scipy.integrate import solve_ivp

# font required by PUP:
plt.rcParams['font.family'] = 'Myriad Pro'

########################################################################
def q(DeltaT,DeltaS):
    ########################################################################
    # MOC transport in m^3/sec as function of temperature and salinity
    # difference between the boxes
    flow=k*(-alpha*DeltaT+beta*DeltaS);
    return flow

########################################################################
def steady_states(Fs,X):
    ########################################################################
    # 3 steady solutions for Y:
    y1=X/2-0.5*np.sqrt(X**2+4*beta*Fs/k)
    if X**2-4*beta*Fs/k>0:
        y2=X/2+0.5*np.sqrt(X**2-4*beta*Fs/k)
        y3=X/2-0.5*np.sqrt(X**2-4*beta*Fs/k)
    else:
        y2=np.nan
        y3=np.nan
    Y=np.array([y1,y2,y3])
    return Y

########################################################################
def rhs_S(time,DeltaS,time_max,DeltaT,hysteresis):
    ########################################################################
    # right hand side of the equation for d/dt(DeltaS):
    # Input: q in m^3/sec; FW is total salt flux into box
    rhs=-(2*np.abs(q(DeltaT,DeltaS))*DeltaS+2*Fs_func(time,time_max,hysteresis));
    return rhs/V

########################################################################
def Fs_func(time,time_max,hysteresis):
    ########################################################################
    # total surface salt flux into northern box

    ########################################################################
    # Specify maximum and minimum of freshwater forcing range during the
    # run in m/year:
    FW_min=-0.1;
    FW_max=5;
    ########################################################################
    if hysteresis:
        if time < time_max/2:
            flux=FW_min+(FW_max-FW_min) \
                *time/(time_max/2);
        else:
            flux=FW_max-(FW_max-FW_min) \
                *(time-(time_max/2))/(time_max/2);
    else:
        flux=2;

    # convert to total salt flux:
    return flux*area*S0/year


########################################################################
## Main program:
########################################################################

# set parameters:
meter=1;
year=365*24*3600;
S0=35.0;
DeltaT=-10;
time_max=10000*year;
alpha=0.2; # (kg/m^3) K^-1
beta=0.8; # (kg/m^3) ppt^-1
k=10e9 # (m^3/s)/(kg/m^3) #  3e-8; # sec^-1
# area of each of the ocean boxes, rediculous value, needed to get
# all other results at the right order of magnitude:
area=(50000.0e3)**2
depth=4000
V=area*depth # volume of each of the ocean boxes
Sv=1.e9 # m^3/sec

print("finding steady states, including unstable ones, as function of F_s...")
# -----------------------------------------------------------------------------
Fs_range=np.arange(0,5,0.01)*area*S0/year
DeltaS_steady=np.zeros((3,len(Fs_range)))
q_steady=np.zeros((3,len(Fs_range)))

# test:
Fs=Fs_func(0.0,0.0,False)
y=steady_states(Fs,alpha*DeltaT)

i=0
for Fs in Fs_range:
    y=steady_states(Fs,alpha*DeltaT)
    DeltaS_steady[:,i]=y/beta
    for j in range(0,3):
        q_steady[j,i]=q(DeltaT,DeltaS_steady[j,i])
    i=i+1

print("calculating dDeltaS/dt as function of DeltaS for stability analysis...")
# ------------------------------------------------------------------------------
hysteresis=False
time=0
DeltaS_range=np.arange(-3,0,0.01)
rhs=np.zeros(len(DeltaS_range))
for i in range(0,len(DeltaS_range)):
    DeltaS = DeltaS_range[i]
    rhs[i]=rhs_S(time,DeltaS,time_max,DeltaT,hysteresis)
    i=i+1
rhs=rhs/np.std(rhs)
    
print("doing a hysteresis run...")
# --------------------------------
hysteresis=True;
y0=[0]
teval=np.arange(0,time_max,time_max/1000)
tspan=(teval[0],teval[-1])
sol = solve_ivp(fun=lambda time,DeltaS: rhs_S(time,DeltaS,time_max,DeltaT,hysteresis) \
                ,vectorized=False,y0=y0,t_span=tspan,t_eval=teval)
T=sol.t
DeltaS=sol.y

hysteresis=True;
FWplot=np.zeros(len(T))
qplot=np.zeros(len(T))
i=0;
for t in T:
    FWplot[i]=Fs_func(t,time_max,hysteresis);
    qplot[i]=q(DeltaT,DeltaS[0,i]);
    i=i+1;
    
N=len(qplot); N2=int(np.floor(N/2));


########################################################################
## plots:
########################################################################
print("plotting...")
plt.figure(1,figsize=(8,4))

Fs_to_m_per_year=S0*area/year
plt.subplot(2,3,1)
plt.plot(Fs_range/Fs_to_m_per_year,DeltaS_steady[0,:],'r.',markersize=1)
plt.plot(Fs_range/Fs_to_m_per_year,DeltaS_steady[1,:],'g.',markersize=1)
plt.plot(Fs_range/Fs_to_m_per_year,DeltaS_steady[2,:],'b.',markersize=1)
plt.plot(Fs_range/Fs_to_m_per_year,0*Fs_range,'k--',dashes=(10, 5),linewidth=0.5)
plt.title('(a) Steady states',loc="left")
plt.xlabel('$F_s$ (m/year)');
plt.ylabel('$\Delta S$ (ppt)');
plt.xlim([min(Fs_range/Fs_to_m_per_year), max(Fs_range/Fs_to_m_per_year)])

plt.subplot(2,3,2)
plt.plot(Fs_range/Fs_to_m_per_year,q_steady[0,:]/Sv,'r.',markersize=1)
plt.plot(Fs_range/Fs_to_m_per_year,q_steady[1,:]/Sv,'g.',markersize=1)
plt.plot(Fs_range/Fs_to_m_per_year,q_steady[2,:]/Sv,'b.',markersize=1)
plt.plot(Fs_range/Fs_to_m_per_year,0*Fs_range,'k--',dashes=(10, 5),linewidth=0.5)
plt.title('(b) Steady states',loc="left")
plt.xlabel('$F_s$ (m/year)');
plt.ylabel('$q$ (Sv)');
plt.xlim([min(Fs_range/Fs_to_m_per_year), max(Fs_range/Fs_to_m_per_year)])

plt.subplot(2,3,3)
plt.plot(DeltaS_range,rhs,'k-',lw=2)
plt.plot(DeltaS_range,rhs*0,'k--',dashes=(10, 5),lw=0.5)
# superimpose color markers of the 3 solutions
Fs=Fs_func(0.0,0.0,False)
yy=steady_states(Fs,alpha*DeltaT)/beta
plt.plot(yy[0],0,'ro',markersize=10)
plt.plot(yy[1],0,'go',markersize=10)
plt.plot(yy[2],0,'bo',markersize=10,fillstyle='none')
plt.title('(c) Stability',loc="left")
plt.xlabel('$\Delta S$ (ppt)');
plt.ylabel('$d\Delta S/dt$');

plt.subplot(2,3,4)
plt.plot(T[:N2]/year/1000,FWplot[:N2]/Fs_to_m_per_year,'b-',markersize=1)
plt.plot(T[N2+1:]/year/1000,FWplot[N2+1:]/Fs_to_m_per_year,'r-',markersize=1)
plt.plot(T/year/1000,0*T,'k--',dashes=(10, 5),linewidth=0.5)
plt.xlabel('Time (kyr)');
plt.ylabel('$F_s$ (m/yr)');
plt.title('(d) $F_s$ for hysteresis run',loc="left");

plt.subplot(2,3,5)
plt.plot(T[:N2]/year/1000,qplot[:N2]/Sv,'b-',markersize=1)
plt.plot(T[N2+1:]/year/1000,qplot[N2+1:]/Sv,'r-',markersize=1)
plt.plot(T/year/1000,T*0,'k--',dashes=(10, 5),lw=0.5)
plt.xlabel('Time (kyr)');
plt.ylabel('MOC, $q$ (Sv)');
plt.title('(e) MOC transport, $q$',loc="left");

plt.subplot(2,3,6)
plt.plot(FWplot[:N2]/Fs_to_m_per_year,qplot[:N2]/Sv,'b-',markersize=1)
plt.plot(FWplot[N2+1:]/Fs_to_m_per_year,qplot[N2+1:]/Sv,'r-',markersize=1)
plt.title('(f) $q$ hysteresis',loc="left");
plt.xlabel('$F_s$ (m/yr)');
plt.ylabel('$q$ (Sv)');


plt.tight_layout()
plt.savefig("Figures/ocean-circulation-stommel-results.pdf")
plt.show()