From 249d243c3765eaf50b55b4342a41d98fbce8d66d Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Tue, 3 Dec 2019 15:58:12 -0500 Subject: [PATCH 01/28] Initial skeleton of model driver --- src/model.py | 48 +++++++++++++++++++++++++++++++++++++ src/pymoc/modules/module.py | 14 +++++++++++ 2 files changed, 62 insertions(+) create mode 100644 src/model.py create mode 100644 src/pymoc/modules/module.py diff --git a/src/model.py b/src/model.py new file mode 100644 index 0000000..9179540 --- /dev/null +++ b/src/model.py @@ -0,0 +1,48 @@ +from pymoc.modules import Module + +class Model(object): + def __init__(self): + self.south_module = None + self.north_module = None + + def keys(self): + keys = [] + module = south_module + while module: + keys.append(module.key) + module = module.north + return keys + + def modules(self): + modules = [] + module = south_module + while module: + modules.append(module) + module = module.north + return modules + + def get_module(self, key): + if not key: + return None + + module = south_module + while module and module.key != key: + module = module.north + if module.key == key: + return module + return None + + def add_module(self, module, name, north_key=None, south_key=None): + module = Module( + module, + name, + north=self.get_module(north_key), + south=self.get_module(south_key), + ) + + if module.north == self.south_module: + self.south_module = module + if module.south == self.north_module: + self.north_module = module + + def run(self): diff --git a/src/pymoc/modules/module.py b/src/pymoc/modules/module.py new file mode 100644 index 0000000..5e57826 --- /dev/null +++ b/src/pymoc/modules/module.py @@ -0,0 +1,14 @@ +class Module(object): + def __init__( + self, + module, + name, + north = None, + south = None, + ): + self.module = module + self.name = name + self.north = north + self.south = south + self.key = self.name + From 0b45d1edac967b5f46a156827347a255a8aade71 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Wed, 4 Dec 2019 17:51:13 -0500 Subject: [PATCH 02/28] Flesh out the (non-functional) model prototype. --- src/model.py | 34 +++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/src/model.py b/src/model.py index 9179540..f1a8d83 100644 --- a/src/model.py +++ b/src/model.py @@ -1,5 +1,6 @@ from pymoc.modules import Module + class Model(object): def __init__(self): self.south_module = None @@ -34,10 +35,10 @@ def get_module(self, key): def add_module(self, module, name, north_key=None, south_key=None): module = Module( - module, - name, - north=self.get_module(north_key), - south=self.get_module(south_key), + module, + name, + north=self.get_module(north_key), + south=self.get_module(south_key), ) if module.north == self.south_module: @@ -45,4 +46,27 @@ def add_module(self, module, name, north_key=None, south_key=None): if module.south == self.north_module: self.north_module = module - def run(self): + def get_modules_by_type(self, module_type): + modules = [] + module = south_module + while module: + if module.module_type == module_type: + modules.append(module) + module = module.north + return modules + + def run(self, dt, steps): + for i in range(0, steps): + self.timestep(dt) + + def timestep(self, dt): + basins = self.get_modules_by_type('basin') + couplers = self.get_modules_by_type('coupler') + for basin in basins: + wAn = basin.north.module.PsiS() if basin.north else 0 + wAs = basin.south.module.PsiN() if basin.south else 0 + wA = (wAn-wAs) * 1e6 + basin.timestep(wA=wA, dt=dt) + for coupler in couplers: + coupler.module.update(b_south=coupler.south, b_north=coupler.north) + coupler.module.solve() From 4bfb32e6e748e5f59c9f2e9a2d39aaccd6352e7f Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 6 Dec 2019 07:40:34 -0500 Subject: [PATCH 03/28] Extract most of the timestep functionality into the module class --- src/model.py | 8 ++--- src/pymoc/modules/SO_ML.py | 2 ++ src/pymoc/modules/column.py | 2 ++ src/pymoc/modules/equi_column.py | 1 + src/pymoc/modules/module.py | 53 +++++++++++++++++++++++++++--- src/pymoc/modules/psi_SO.py | 1 + src/pymoc/modules/psi_thermwind.py | 2 ++ 7 files changed, 58 insertions(+), 11 deletions(-) diff --git a/src/model.py b/src/model.py index f1a8d83..3e203be 100644 --- a/src/model.py +++ b/src/model.py @@ -63,10 +63,6 @@ def timestep(self, dt): basins = self.get_modules_by_type('basin') couplers = self.get_modules_by_type('coupler') for basin in basins: - wAn = basin.north.module.PsiS() if basin.north else 0 - wAs = basin.south.module.PsiN() if basin.south else 0 - wA = (wAn-wAs) * 1e6 - basin.timestep(wA=wA, dt=dt) + basin.timestep(dt=dt) for coupler in couplers: - coupler.module.update(b_south=coupler.south, b_north=coupler.north) - coupler.module.solve() + coupler.timestep() diff --git a/src/pymoc/modules/SO_ML.py b/src/pymoc/modules/SO_ML.py index dc5964b..2bdfd98 100644 --- a/src/pymoc/modules/SO_ML.py +++ b/src/pymoc/modules/SO_ML.py @@ -273,6 +273,8 @@ def advdiff(self, b_basin, Psi_b, dt): # (preferable e.g. for computation of streamfunction) self.set_boundary_conditions(b_basin, Psi_b) + self.module_type = 'basin' + def timestep(self, b_basin=None, Psi_b=None, dt=1.): r""" Integrate the mixed layer buoyancy profile for one timestep. diff --git a/src/pymoc/modules/column.py b/src/pymoc/modules/column.py index 622498e..e98374a 100644 --- a/src/pymoc/modules/column.py +++ b/src/pymoc/modules/column.py @@ -71,6 +71,8 @@ def __init__( else: self.bz = 0. * z # notice that this is just for initialization of ode solver + self.module_type = 'basin' + def Akappa(self, z): r""" Compute the area integrated diffusivity :math:`A\kappa` diff --git a/src/pymoc/modules/equi_column.py b/src/pymoc/modules/equi_column.py index afe4352..49fd1ed 100644 --- a/src/pymoc/modules/equi_column.py +++ b/src/pymoc/modules/equi_column.py @@ -358,6 +358,7 @@ def bc(self, ya, yb, p=None): raise TypeError( 'Must provide a p array if column does not have an H value' ) + self.module_type = 'basin' def ode(self, z, y, p=None): r""" diff --git a/src/pymoc/modules/module.py b/src/pymoc/modules/module.py index 5e57826..5ad9208 100644 --- a/src/pymoc/modules/module.py +++ b/src/pymoc/modules/module.py @@ -1,10 +1,10 @@ class Module(object): def __init__( - self, - module, - name, - north = None, - south = None, + self, + module, + name, + north=None, + south=None, ): self.module = module self.name = name @@ -12,3 +12,46 @@ def __init__( self.south = south self.key = self.name + @property + def module_type(self): + return self.module.module_type + + @property + def psi_north(self): + if self.north and self.north.module_type == 'coupler': + coupler = self.north.module + if hasattr(coupler, 'Psibz') and callable(coupler, 'Psibz'): + [Psi, _] = coupler.Psibz() + return Psi + else: + return None + return None + + @property + def psi_south(self): + if self.south and self.south.module_type == 'coupler': + coupler = self.south.module + if hasattr(coupler, 'Psibz') and callable(coupler, 'Psibz'): + [_, Psi] = coupler.Psibz() + return Psi + else: + return coupler.Psi + return None + + def timestep(self, dt=None): + if self.module_type == 'basin': + wAn = self.psi_north + wAs = self.psi_south + wA = (wAn-wAs) * 1e6 + self.module.timestep(wA=wA, dt=dt) + if self.north: + self.north.south = self + if self.south: + self.south.north = self + elif self.module_type == 'coupler': + self.module.update(b_south=self.south, b_north=self.north) + self.module.solve() + if self.north: + self.north.south = self + if self.south: + self.south.north = self diff --git a/src/pymoc/modules/psi_SO.py b/src/pymoc/modules/psi_SO.py index 8ee342b..19913ed 100644 --- a/src/pymoc/modules/psi_SO.py +++ b/src/pymoc/modules/psi_SO.py @@ -353,6 +353,7 @@ def solve(self): # only for BC. To avoid random fluctuation in bottom Psi, we here simply # set it to value of last gridpoint above self.Psi[0] = self.Psi[1] + self.module_type = 'coupler' def update(self, b=None, bs=None): r""" diff --git a/src/pymoc/modules/psi_thermwind.py b/src/pymoc/modules/psi_thermwind.py index c965219..9c981a5 100644 --- a/src/pymoc/modules/psi_thermwind.py +++ b/src/pymoc/modules/psi_thermwind.py @@ -69,6 +69,8 @@ def __init__( else: self.sol_init = sol_init + self.module_type = 'coupler' + def bc(self, ya, yb): r""" Calculate the residuals of boundary conditions for the thermal wind closure From a43499a96b842a7008d33bd84741323fbe243e7a Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Mon, 16 Dec 2019 18:57:57 -0500 Subject: [PATCH 04/28] improve model performance --- .codecov.yml | 18 -- examples/example_twocol_plusSO.py | 15 +- examples/example_twocol_plusSO_model.py | 246 ++++++++++++++++++++++++ src/model.py | 76 +++++--- src/pymoc/__init__.py | 1 + src/pymoc/modules/__init__.py | 1 + src/pymoc/modules/module.py | 69 ++++--- src/pymoc/modules/psi_SO.py | 4 +- src/pymoc/modules/psi_thermwind.py | 1 + 9 files changed, 363 insertions(+), 68 deletions(-) delete mode 100644 .codecov.yml create mode 100644 examples/example_twocol_plusSO_model.py diff --git a/.codecov.yml b/.codecov.yml deleted file mode 100644 index b8768f5..0000000 --- a/.codecov.yml +++ /dev/null @@ -1,18 +0,0 @@ -# Team Yaml -coverage: - round: down - precision: 2 - -# Repository Yaml -coverage: - round: down - range: 40..80 - - ignore: - - "examples/**/*.py" - - "tests/**/*.py" - - "setup.py" - - "pymoc/plotting/*" - -codecov: - max_report_age: off diff --git a/examples/example_twocol_plusSO.py b/examples/example_twocol_plusSO.py index 3945374..ce19162 100644 --- a/examples/example_twocol_plusSO.py +++ b/examples/example_twocol_plusSO.py @@ -13,6 +13,7 @@ import numpy as np from matplotlib import pyplot as plt from scipy.interpolate import interp1d +from datetime import datetime # boundary conditions: bs = 0.03 @@ -54,8 +55,10 @@ # Initial conditions for buoyancy profile in the basin def b_basin(z): return bs * np.exp(z / 300.) + + def b_north(z): - return 1e-3*bs * np.exp(z / 300.) + return 1e-3 * bs * np.exp(z / 300.) # create N.A. overturning model instance @@ -82,7 +85,9 @@ def b_north(z): # create adv-diff column model instance for basin basin = Column(z=z, kappa=kappa, Area=A_basin, b=b_basin, bs=bs, bbot=bmin) # create adv-diff column model instance for basin -north = Column(z=z, kappa=kappa, Area=A_north, b=b_north, bs=bs_north, bbot=bmin) +north = Column( + z=z, kappa=kappa, Area=A_north, b=b_north, bs=bs_north, bbot=bmin +) # Create figure: fig = plt.figure(figsize=(6, 10)) @@ -101,6 +106,8 @@ def b_north(z): #wAb=(AMOC.Psi-SO.Psi)*1e6 #wAN=-AMOC.Psi*1e6 # using isopycnal overturning: + print(str(ii) + '/' + str(total_iters)) + print(datetime.now()) wAb = (Psi_iso_b - SO.Psi) * 1e6 wAN = -Psi_iso_n * 1e6 basin.timestep(wA=wAb, dt=dt) @@ -114,11 +121,14 @@ def b_north(z): SO.solve() if ii % plot_iters == 0: # Plot current state: + ax1.cla() + ax2.cla() ax1.plot(AMOC.Psi, AMOC.z, linewidth=0.5, color='r') ax1.plot(SO.Psi, SO.z, linewidth=0.5, color='m') ax2.plot(basin.b, basin.z, linewidth=0.5, color='b') ax2.plot(north.b, north.z, linewidth=0.5, color='c') plt.pause(0.01) + print(datetime.now()) # Plot final results over time-iteration plot: ax1.plot(AMOC.Psi, AMOC.z, linewidth=2, color='r') @@ -240,3 +250,4 @@ def b_north(z): ax1.set_xlabel('y [km]') ax1.set_ylabel('Depth [m]') ax1.set_title('Depth-averaged Overturning') +plt.show() diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py new file mode 100644 index 0000000..59c49ea --- /dev/null +++ b/examples/example_twocol_plusSO_model.py @@ -0,0 +1,246 @@ +''' +This script shows an example of a "two column" model for the +overturning circulation in a basin connected to a channel in the south. +The first column represents the basin, while the second column represents +the northern sinking region. The overtunning circulation is computed at +the northern end of the basin (at the interface to the northern sinking region) +and at the southern end of the basin (at the interface to the channel). +The parameters chosen here follow more or less the "control" experiment of Nikurashin +and Vallis (2012, JPO). +''' +from pymoc import Model +from pymoc.modules import Psi_Thermwind, Psi_SO, Column +from pymoc.plotting import Interpolate_channel +import numpy as np +from matplotlib import pyplot as plt +from scipy.interpolate import interp1d +import cProfile + +# boundary conditions: +bs = 0.03 +bs_north = 0.005 +bmin = 0.0 + +# S.O. surface boundary conditions and grid: +l = 2.e6 +y = np.asarray(np.linspace(0, l, 40)) +#Although the model can emulate the effect of a meridionally varying wind-stress, +# that feature is in "beta" and we are here for simplicity using a constant +# wind-stress (approximately the average over the channel in NV12): +tau = 0.13 +# from what I can infer, NV12 use an approximately quadratic surface temperature +# profile in the channel for their SAMBUCA simulations: +bs_SO = (bs-bmin) * (y / y[-1])**2 + bmin + +A_basin = 6e13 #area of the basin +A_north = A_basin / 50. #area of northern sinking region + +# time-stepping parameters: +dt = 86400 * 30 # time-step for vert. adv. diff. calc. +MOC_up_iters = int( + np.floor(2. * 360 * 86400 / dt) +) # multiplier for MOC time-step (MOC is updated every MOC_up_iters time steps) +plot_iters = int( + np.ceil(300 * 360 * 86400 / dt) +) # plotting frequency (in iterations) +total_iters = int( + np.ceil(3000 * 360 * 86400 / dt) +) # total number of timesteps + +kappa = 2e-5 + +# create vertical grid: +z = np.asarray(np.linspace(-4000, 0, 80)) + + +# Initial conditions for buoyancy profile in the basin +def b_basin(z): + return bs * np.exp(z / 300.) + + +def b_north(z): + return 1e-3 * bs * np.exp(z / 300.) + + +# create N.A. overturning model instance +AMOC = Psi_Thermwind(z=z, b1=b_basin, b2=b_north, f=1e-4) +SO = Psi_SO( + z=z, + y=y, + b=b_basin(z), + bs=bs_SO, + tau=tau, + f=1e-4, + L=5e6, + KGM=1000., + c=0.1, + bvp_with_Ek=False +) +# create adv-diff column model instance for basin +basin = Column(z=z, kappa=kappa, Area=A_basin, b=b_basin, bs=bs, bbot=bmin) +# create adv-diff column model instance for basin +north = Column( + z=z, kappa=kappa, Area=A_north, b=b_north, bs=bs_north, bbot=bmin +) + +model = Model() +model.add_module(SO, 'Psi SO') +model.add_module(basin, 'Atlantic Basin', south_key='psi_so') +model.add_module(AMOC, 'AMOC', south_key='atlantic_basin') +model.add_module(north, 'North Atlantic', south_key='amoc', do_conv=True) + +fig = plt.figure(figsize=(6, 10)) +ax1 = fig.add_subplot(111) +ax2 = ax1.twiny() +plt.ylim((-4e3, 0)) +ax1.set_xlim((-10, 15)) +ax2.set_xlim((-0.02, 0.03)) +ax1.set_xlabel('$\Psi$', fontsize=14) +ax2.set_xlabel('b', fontsize=14) + +AMOC = model.get_module('amoc').module +SO = model.get_module('psi_so').module +basin = model.get_module('atlantic_basin').module +north = model.get_module('north_atlantic').module + +for snapshot in model.run( + basin_dt=dt, + coupler_dt=MOC_up_iters, + steps=total_iters, + snapshot_start=plot_iters, + snapshot_interval=plot_iters +): + ax1.cla() + ax2.cla() + ax1.plot(AMOC.Psi, AMOC.z, linewidth=0.5, color='r') + ax1.plot(SO.Psi, SO.z, linewidth=0.5, color='m') + ax2.plot(basin.b, basin.z, linewidth=0.5, color='b') + ax2.plot(north.b, north.z, linewidth=0.5, color='c') + plt.pause(0.01) + +plt.show() +# AMOC = model.get_module('amoc').module +# SO = model.get_module('psi_so').module +# basin = model.get_module('atlantic_basin').module +# north = model.get_module('north_atlantic').module +# # Plot final results over time-iteration plot: +ax1.plot(AMOC.Psi, AMOC.z, linewidth=2, color='r') +ax1.plot(SO.Psi, SO.z, linewidth=2, color='m') +ax1.plot(SO.Psi_Ek, SO.z, linewidth=1, color='m', linestyle='--') +ax1.plot(SO.Psi_GM, SO.z, linewidth=1, color='m', linestyle=':') +ax2.plot(basin.b, basin.z, linewidth=2, color='b') +ax2.plot(north.b, basin.z, linewidth=2, color='c') +ax1.plot(0. * AMOC.z, AMOC.z, linewidth=0.5, color='k') + +# #============================================================================== +# # Everything below is just to make fancy 2D plots: +# # This is mostly an exercise in interpolation, but useful to visualize solutions + +# # first interpolate buoyancy in channel along constant-slope isopycnals: +bint = Interpolate_channel(y=y, z=z, bs=bs_SO, bn=basin.b) +bsouth = bint.gridit() +# buoyancy in the basin is all the same: +ybasin = np.linspace(200., 12000., 60) +bbasin = np.tile(basin.b, (len(ybasin), 1)) +# finally, interpolate buoyancy in the north: +ynorth = np.linspace(12100., 13000., 10) +bnorth = np.zeros((len(ynorth), len(z))) +for iz in range(len(z)): + bnorth[:, iz] = interp1d([11900., 12000., 12900., 13000.], [ + basin.b[iz], basin.b[iz], north.b[iz], north.b[iz] + ], + kind='quadratic')(ynorth) +# now stick it all together: +ynew = np.concatenate(((y-l) / 1e3, ybasin, ynorth)) +bnew = np.concatenate((bsouth, bbasin, bnorth)) + +# Evaluate isopycnal overturning streamfunction at all latitudes: +psiarray_iso = np.zeros((len(ynew), len(z))) +psiarray_z = np.zeros((len(ynew), len(z))) +for iy in range(1, len(y)): + # in the channel, interpolate SO.Psi onto local isopycnal depth: + psiarray_iso[iy, :] = np.interp(bnew[iy, :], basin.b, SO.Psi) + psiarray_z[iy, :] = psiarray_iso[iy, :] +for iy in range(len(y), len(y) + len(ybasin)): + # in the basin, linearly interpolate between Psi_SO and Psi_AMOC: + psiarray_iso[iy, :] = ( + ynew[iy] * AMOC.Psibz()[0] + (10000. - ynew[iy]) * SO.Psi + ) / 10000. + psiarray_z[iy, : + ] = (ynew[iy] * AMOC.Psi + (10000. - ynew[iy]) * SO.Psi) / 10000. +for iy in range(len(y) + len(ybasin), len(ynew)): + # in the north, interpolate AMOC.psib to local isopycnal depth: + psiarray_iso[iy, :] = np.interp(bnew[iy, :], AMOC.bgrid, AMOC.Psib()) + psiarray_z[iy, :] = ((13000. - ynew[iy]) * AMOC.Psi) / 1000. +psiarray_iso[-1, :] = 0. + +blevs = np.arange(0.001, 0.03, 0.002) +# Notice that the contour level for Psi is 0.5 SV for negative values +# and 1 SV for positive values. This appers to be what was used in NV2012. +plevs = np.concatenate((np.arange(-15., -0.4, 0.5), np.arange(1., 16., 1.))) + +# plot isopycnal overturning: +fig = plt.figure(figsize=(10, 5)) +ax1 = fig.add_subplot(111) +CS = ax1.contour( + ynew, + z, + bnew.transpose(), + levels=blevs, + colors='k', + linewidths=1.0, + linestyles='solid' +) +ax1.clabel(CS, fontsize=10) +ax1.contour( + ynew, + z, + psiarray_iso.transpose(), + levels=plevs, + colors='k', + linewidths=0.5 +) +CS = ax1.contourf( + ynew, + z, + psiarray_iso.transpose(), + levels=plevs, + cmap=plt.cm.bwr, + vmin=-15, + vmax=15 +) +#fig.colorbar(CS, ticks=plevs[0::2], orientation='vertical') +ax1.set_xlabel('y [km]') +ax1.set_ylabel('Depth [m]') +ax1.set_title('Isopycnal Overturning') + +# plot z-coord. overturning: +fig = plt.figure(figsize=(10, 5)) +ax1 = fig.add_subplot(111) +CS = ax1.contour( + ynew, + z, + bnew.transpose(), + levels=blevs, + colors='k', + linewidths=1.0, + linestyles='solid' +) +ax1.clabel(CS, fontsize=10) +ax1.contour( + ynew, z, psiarray_z.transpose(), levels=plevs, colors='k', linewidths=0.5 +) +CS = ax1.contourf( + ynew, + z, + psiarray_z.transpose(), + levels=plevs, + cmap=plt.cm.bwr, + vmin=-15, + vmax=15 +) +#fig.colorbar(CS, ticks=plevs[0::2], orientation='vertical') +ax1.set_xlabel('y [km]') +ax1.set_ylabel('Depth [m]') +ax1.set_title('Depth-averaged Overturning') +plt.show() diff --git a/src/model.py b/src/model.py index 3e203be..d67511d 100644 --- a/src/model.py +++ b/src/model.py @@ -5,10 +5,13 @@ class Model(object): def __init__(self): self.south_module = None self.north_module = None + self.basins = [] + self.couplers = [] + self._modules = {} def keys(self): keys = [] - module = south_module + module = self.south_module while module: keys.append(module.key) module = module.north @@ -16,7 +19,7 @@ def keys(self): def modules(self): modules = [] - module = south_module + module = self.south_module while module: modules.append(module) module = module.north @@ -26,43 +29,70 @@ def get_module(self, key): if not key: return None - module = south_module - while module and module.key != key: - module = module.north - if module.key == key: - return module - return None + return self._modules[key] + # module = self.south_module + # while module and module.key != key: + # module = module.north + # if module and module.key == key: + # return module + # return None - def add_module(self, module, name, north_key=None, south_key=None): + def add_module( + self, module, name, north_key=None, south_key=None, do_conv=False + ): + north = self.get_module(north_key) + south = self.get_module(south_key) module = Module( module, name, - north=self.get_module(north_key), - south=self.get_module(south_key), + north=north, + south=south, + do_conv=do_conv, ) + self._modules[module.key] = module - if module.north == self.south_module: + if north == self.south_module: self.south_module = module - if module.south == self.north_module: + if south == self.north_module: self.north_module = module + if module.module_type == 'basin': + self.basins.append(module) + elif module.module_type == 'coupler': + self.couplers.append(module) + def get_modules_by_type(self, module_type): modules = [] - module = south_module + module = self.south_module while module: if module.module_type == module_type: modules.append(module) module = module.north return modules - def run(self, dt, steps): + def run( + self, + steps, + basin_dt, + coupler_dt, + snapshot_start=None, + snapshot_interval=None + ): for i in range(0, steps): - self.timestep(dt) + print(str(i) + '/' + str(steps)) + self.timestep(i, basin_dt, coupler_dt=coupler_dt) + if self.snapshot(i, snapshot_start, snapshot_interval): + yield i + return + + def timestep(self, step, basin_dt, coupler_dt=0): + for basin in self.basins: + basin.timestep(dt=basin_dt) + if step % coupler_dt == 0: + for coupler in self.couplers: + coupler.timestep() - def timestep(self, dt): - basins = self.get_modules_by_type('basin') - couplers = self.get_modules_by_type('coupler') - for basin in basins: - basin.timestep(dt=dt) - for coupler in couplers: - coupler.timestep() + def snapshot(self, step, snapshot_start, snapshot_interval): + return snapshot_start is not None and snapshot_interval is not None and step >= snapshot_start and ( + step-snapshot_start + ) % snapshot_interval == 0 diff --git a/src/pymoc/__init__.py b/src/pymoc/__init__.py index ac05284..d576a8a 100644 --- a/src/pymoc/__init__.py +++ b/src/pymoc/__init__.py @@ -2,3 +2,4 @@ import pymoc.modules import pymoc.utils import pymoc.plotting +from model import Model \ No newline at end of file diff --git a/src/pymoc/modules/__init__.py b/src/pymoc/modules/__init__.py index 1eb3639..e9d367c 100644 --- a/src/pymoc/modules/__init__.py +++ b/src/pymoc/modules/__init__.py @@ -4,3 +4,4 @@ from .psi_SO import Psi_SO from .psi_thermwind import Psi_Thermwind from .SO_ML import SO_ML +from .module import Module diff --git a/src/pymoc/modules/module.py b/src/pymoc/modules/module.py index 5ad9208..2eef692 100644 --- a/src/pymoc/modules/module.py +++ b/src/pymoc/modules/module.py @@ -1,3 +1,7 @@ +from datetime import datetime +from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column + + class Module(object): def __init__( self, @@ -5,12 +9,23 @@ def __init__( name, north=None, south=None, + do_conv=False, ): + if north: + north.south = self + if south: + south.north = self self.module = module self.name = name self.north = north self.south = south - self.key = self.name + self.do_conv = do_conv + self.key = name.replace(' ', '_').lower().strip('_') + self.do_psi_bz = hasattr(module, 'Psibz') and callable(module.Psibz) + self.b_type = 'b' if isinstance(self.module, Column) or isinstance( + self.module, Equi_Column + ) else 'bs' + # self.key = ''.join(['_'+i.lower() if i.isupper() else i for i in name]).lstrip('_').replace(' ', '') @property def module_type(self): @@ -18,40 +33,46 @@ def module_type(self): @property def psi_north(self): - if self.north and self.north.module_type == 'coupler': - coupler = self.north.module - if hasattr(coupler, 'Psibz') and callable(coupler, 'Psibz'): - [Psi, _] = coupler.Psibz() + coupler = self.north + if coupler and coupler.module_type == 'coupler': + if coupler.do_psi_bz: + [Psi, _] = coupler.module.Psibz() return Psi else: - return None - return None + return 0 + return 0 @property def psi_south(self): - if self.south and self.south.module_type == 'coupler': - coupler = self.south.module - if hasattr(coupler, 'Psibz') and callable(coupler, 'Psibz'): - [_, Psi] = coupler.Psibz() + coupler = self.south + if coupler and coupler.module_type == 'coupler': + if coupler.do_psi_bz: + [_, Psi] = coupler.module.Psibz() return Psi else: - return coupler.Psi - return None + return coupler.module.Psi + return 0 def timestep(self, dt=None): + module = self.module if self.module_type == 'basin': wAn = self.psi_north wAs = self.psi_south wA = (wAn-wAs) * 1e6 - self.module.timestep(wA=wA, dt=dt) - if self.north: - self.north.south = self - if self.south: - self.south.north = self + module.timestep(wA=wA, dt=dt, do_conv=self.do_conv) elif self.module_type == 'coupler': - self.module.update(b_south=self.south, b_north=self.north) - self.module.solve() - if self.north: - self.north.south = self - if self.south: - self.south.north = self + north = self.north + south = self.south + if self.do_psi_bz: + module.update(b2=north and north.b, b1=south and south.b) + else: + module.update(b=north and north.b) + module.solve() + + @property + def b(self): + return self[self.b_type] + if isinstance(self.module, Column) or isinstance(self.module, Equi_Column): + return self.module.b + if isinstance(self.module, SO_ML): + return self.module.bs diff --git a/src/pymoc/modules/psi_SO.py b/src/pymoc/modules/psi_SO.py index 19913ed..4c03b3c 100644 --- a/src/pymoc/modules/psi_SO.py +++ b/src/pymoc/modules/psi_SO.py @@ -103,6 +103,9 @@ def __init__( self.Htaperbot = Htaperbot self.smax = smax + self.Psi = np.zeros(len(self.z)) + self.module_type = 'coupler' + def ys(self, b): r""" Inversion function of :math:`bs\left(y\right)`. This is equivalent to the outcopping @@ -353,7 +356,6 @@ def solve(self): # only for BC. To avoid random fluctuation in bottom Psi, we here simply # set it to value of last gridpoint above self.Psi[0] = self.Psi[1] - self.module_type = 'coupler' def update(self, b=None, bs=None): r""" diff --git a/src/pymoc/modules/psi_thermwind.py b/src/pymoc/modules/psi_thermwind.py index 9c981a5..2047a91 100644 --- a/src/pymoc/modules/psi_thermwind.py +++ b/src/pymoc/modules/psi_thermwind.py @@ -69,6 +69,7 @@ def __init__( else: self.sol_init = sol_init + self.Psi = np.zeros(len(self.z)) self.module_type = 'coupler' def bc(self, ya, yb): From 4d0a0ddeaea90a66094d875987cea68b3d780300 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Tue, 17 Dec 2019 07:56:11 -0500 Subject: [PATCH 05/28] fix model installation, performance tweak --- examples/example_twocol_plusSO_model.py | 6 +++--- src/pymoc/__init__.py | 2 +- src/{ => pymoc}/model.py | 0 src/pymoc/modules/module.py | 6 +----- 4 files changed, 5 insertions(+), 9 deletions(-) rename src/{ => pymoc}/model.py (100%) diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index 59c49ea..2dbc61a 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -8,14 +8,14 @@ The parameters chosen here follow more or less the "control" experiment of Nikurashin and Vallis (2012, JPO). ''' -from pymoc import Model +from pymoc import model from pymoc.modules import Psi_Thermwind, Psi_SO, Column from pymoc.plotting import Interpolate_channel import numpy as np from matplotlib import pyplot as plt from scipy.interpolate import interp1d -import cProfile +print(model.Model) # boundary conditions: bs = 0.03 bs_north = 0.005 @@ -83,7 +83,7 @@ def b_north(z): z=z, kappa=kappa, Area=A_north, b=b_north, bs=bs_north, bbot=bmin ) -model = Model() +model = model.Model() model.add_module(SO, 'Psi SO') model.add_module(basin, 'Atlantic Basin', south_key='psi_so') model.add_module(AMOC, 'AMOC', south_key='atlantic_basin') diff --git a/src/pymoc/__init__.py b/src/pymoc/__init__.py index d576a8a..d83ad0e 100644 --- a/src/pymoc/__init__.py +++ b/src/pymoc/__init__.py @@ -2,4 +2,4 @@ import pymoc.modules import pymoc.utils import pymoc.plotting -from model import Model \ No newline at end of file +import pymoc.model \ No newline at end of file diff --git a/src/model.py b/src/pymoc/model.py similarity index 100% rename from src/model.py rename to src/pymoc/model.py diff --git a/src/pymoc/modules/module.py b/src/pymoc/modules/module.py index 2eef692..8c08a93 100644 --- a/src/pymoc/modules/module.py +++ b/src/pymoc/modules/module.py @@ -71,8 +71,4 @@ def timestep(self, dt=None): @property def b(self): - return self[self.b_type] - if isinstance(self.module, Column) or isinstance(self.module, Equi_Column): - return self.module.b - if isinstance(self.module, SO_ML): - return self.module.bs + return getattr(self.module, self.b_type) From c47ee96ddca9cf8b6d218f909ebb11009e82a058 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 7 Feb 2020 11:52:46 -0600 Subject: [PATCH 06/28] Only update Psibz at coupler timestep -Interpolation onto depth space is very expensive -Does not need to happen at every timestep --- examples/example_twocol_plusSO.py | 4 -- examples/example_twocol_plusSO_model.py | 1 - src/pymoc/model.py | 5 +-- src/pymoc/modules/module.py | 53 +++++++++---------------- 4 files changed, 20 insertions(+), 43 deletions(-) diff --git a/examples/example_twocol_plusSO.py b/examples/example_twocol_plusSO.py index ce19162..c427dfb 100644 --- a/examples/example_twocol_plusSO.py +++ b/examples/example_twocol_plusSO.py @@ -13,7 +13,6 @@ import numpy as np from matplotlib import pyplot as plt from scipy.interpolate import interp1d -from datetime import datetime # boundary conditions: bs = 0.03 @@ -106,8 +105,6 @@ def b_north(z): #wAb=(AMOC.Psi-SO.Psi)*1e6 #wAN=-AMOC.Psi*1e6 # using isopycnal overturning: - print(str(ii) + '/' + str(total_iters)) - print(datetime.now()) wAb = (Psi_iso_b - SO.Psi) * 1e6 wAN = -Psi_iso_n * 1e6 basin.timestep(wA=wAb, dt=dt) @@ -128,7 +125,6 @@ def b_north(z): ax2.plot(basin.b, basin.z, linewidth=0.5, color='b') ax2.plot(north.b, north.z, linewidth=0.5, color='c') plt.pause(0.01) - print(datetime.now()) # Plot final results over time-iteration plot: ax1.plot(AMOC.Psi, AMOC.z, linewidth=2, color='r') diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index 2dbc61a..ac6cebd 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -15,7 +15,6 @@ from matplotlib import pyplot as plt from scipy.interpolate import interp1d -print(model.Model) # boundary conditions: bs = 0.03 bs_north = 0.005 diff --git a/src/pymoc/model.py b/src/pymoc/model.py index d67511d..f289656 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -79,7 +79,6 @@ def run( snapshot_interval=None ): for i in range(0, steps): - print(str(i) + '/' + str(steps)) self.timestep(i, basin_dt, coupler_dt=coupler_dt) if self.snapshot(i, snapshot_start, snapshot_interval): yield i @@ -87,10 +86,10 @@ def run( def timestep(self, step, basin_dt, coupler_dt=0): for basin in self.basins: - basin.timestep(dt=basin_dt) + basin.timestep_basin(dt=basin_dt) if step % coupler_dt == 0: for coupler in self.couplers: - coupler.timestep() + coupler.update_coupler() def snapshot(self, step, snapshot_start, snapshot_interval): return snapshot_start is not None and snapshot_interval is not None and step >= snapshot_start and ( diff --git a/src/pymoc/modules/module.py b/src/pymoc/modules/module.py index 8c08a93..cbe7dcf 100644 --- a/src/pymoc/modules/module.py +++ b/src/pymoc/modules/module.py @@ -25,49 +25,32 @@ def __init__( self.b_type = 'b' if isinstance(self.module, Column) or isinstance( self.module, Equi_Column ) else 'bs' + self.psi = [0, 0] # self.key = ''.join(['_'+i.lower() if i.isupper() else i for i in name]).lstrip('_').replace(' ', '') @property def module_type(self): return self.module.module_type - @property - def psi_north(self): - coupler = self.north - if coupler and coupler.module_type == 'coupler': - if coupler.do_psi_bz: - [Psi, _] = coupler.module.Psibz() - return Psi - else: - return 0 - return 0 - - @property - def psi_south(self): - coupler = self.south - if coupler and coupler.module_type == 'coupler': - if coupler.do_psi_bz: - [_, Psi] = coupler.module.Psibz() - return Psi - else: - return coupler.module.Psi - return 0 + def timestep_basin(self, dt=None): + module = self.module + wA = 0.0 + if self.north: + wA += self.north.psi[0] * 1e6 + if self.south: + wA -= self.south.psi[-1] * 1e6 + module.timestep(wA=wA, dt=dt, do_conv=self.do_conv) - def timestep(self, dt=None): + def update_coupler(self, dt=None): module = self.module - if self.module_type == 'basin': - wAn = self.psi_north - wAs = self.psi_south - wA = (wAn-wAs) * 1e6 - module.timestep(wA=wA, dt=dt, do_conv=self.do_conv) - elif self.module_type == 'coupler': - north = self.north - south = self.south - if self.do_psi_bz: - module.update(b2=north and north.b, b1=south and south.b) - else: - module.update(b=north and north.b) - module.solve() + north = self.north + south = self.south + if self.do_psi_bz: + module.update(b2=north and north.b, b1=south and south.b) + else: + module.update(b=north and north.b) + module.solve() + self.psi = module.Psibz() if self.do_psi_bz else [module.Psi] @property def b(self): From b6318c887f0c6dc715972187972d5b95dc6300ee Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 7 Feb 2020 12:09:00 -0600 Subject: [PATCH 07/28] Add add_module methond to model -Rename Module to ModuleWrapper to avoid confustion -Allow user to defer module construction to the wrapper --- examples/example_twocol_plusSO_model.py | 71 +++++++++++++------ src/pymoc/model.py | 21 +++++- src/pymoc/modules/__init__.py | 2 +- .../modules/{module.py => module_wrapper.py} | 2 +- 4 files changed, 70 insertions(+), 26 deletions(-) rename src/pymoc/modules/{module.py => module_wrapper.py} (98%) diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index ac6cebd..b76428f 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -61,33 +61,60 @@ def b_north(z): return 1e-3 * bs * np.exp(z / 300.) -# create N.A. overturning model instance -AMOC = Psi_Thermwind(z=z, b1=b_basin, b2=b_north, f=1e-4) -SO = Psi_SO( - z=z, - y=y, - b=b_basin(z), - bs=bs_SO, - tau=tau, - f=1e-4, - L=5e6, - KGM=1000., - c=0.1, - bvp_with_Ek=False +model = model.Model() +model.new_module( + Psi_SO, { + 'z': z, + 'y': y, + 'b': b_basin(z), + 'bs': bs_SO, + 'tau': tau, + 'f': 1e-4, + 'L': 5e6, + 'KGM': 1000., + 'c': 0.1, + 'bvp_with_Ek': False + }, 'Psi SO' ) # create adv-diff column model instance for basin -basin = Column(z=z, kappa=kappa, Area=A_basin, b=b_basin, bs=bs, bbot=bmin) +model.new_module( + Column, { + 'z': z, + 'kappa': kappa, + 'Area': A_basin, + 'b': b_basin, + 'bs': bs, + 'bbot': bmin + }, + 'Atlantic Basin', + south_key='psi_so' +) +# create N.A. overturning model instance +model.new_module( + Psi_Thermwind, { + 'z': z, + 'b1': b_basin, + 'b2': b_north, + 'f': 1e-4 + }, + 'AMOC', + south_key='atlantic_basin' +) # create adv-diff column model instance for basin -north = Column( - z=z, kappa=kappa, Area=A_north, b=b_north, bs=bs_north, bbot=bmin +model.new_module( + Column, { + 'z': z, + 'kappa': kappa, + 'Area': A_north, + 'b': b_north, + 'bs': bs_north, + 'bbot': bmin + }, + 'North Atlantic', + south_key='amoc', + do_conv=True ) -model = model.Model() -model.add_module(SO, 'Psi SO') -model.add_module(basin, 'Atlantic Basin', south_key='psi_so') -model.add_module(AMOC, 'AMOC', south_key='atlantic_basin') -model.add_module(north, 'North Atlantic', south_key='amoc', do_conv=True) - fig = plt.figure(figsize=(6, 10)) ax1 = fig.add_subplot(111) ax2 = ax1.twiny() diff --git a/src/pymoc/model.py b/src/pymoc/model.py index f289656..4a3dbeb 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -1,4 +1,4 @@ -from pymoc.modules import Module +from pymoc.modules import ModuleWrapper class Model(object): @@ -42,7 +42,7 @@ def add_module( ): north = self.get_module(north_key) south = self.get_module(south_key) - module = Module( + module = ModuleWrapper( module, name, north=north, @@ -61,6 +61,23 @@ def add_module( elif module.module_type == 'coupler': self.couplers.append(module) + def new_module( + self, + module_class, + module_args, + module_name, + north_key=None, + south_key=None, + do_conv=False + ): + self.add_module( + module_class(**module_args), + module_name, + north_key=north_key, + south_key=south_key, + do_conv=do_conv + ) + def get_modules_by_type(self, module_type): modules = [] module = self.south_module diff --git a/src/pymoc/modules/__init__.py b/src/pymoc/modules/__init__.py index e9d367c..0c85360 100644 --- a/src/pymoc/modules/__init__.py +++ b/src/pymoc/modules/__init__.py @@ -4,4 +4,4 @@ from .psi_SO import Psi_SO from .psi_thermwind import Psi_Thermwind from .SO_ML import SO_ML -from .module import Module +from .module_wrapper import ModuleWrapper diff --git a/src/pymoc/modules/module.py b/src/pymoc/modules/module_wrapper.py similarity index 98% rename from src/pymoc/modules/module.py rename to src/pymoc/modules/module_wrapper.py index cbe7dcf..d9617ca 100644 --- a/src/pymoc/modules/module.py +++ b/src/pymoc/modules/module_wrapper.py @@ -2,7 +2,7 @@ from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column -class Module(object): +class ModuleWrapper(object): def __init__( self, module, From 9e9e771beb3f4231004cd36b6e9a0f25b1cd00fb Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 7 Feb 2020 12:13:09 -0600 Subject: [PATCH 08/28] Module -> ModuleWrapper --- src/pymoc/model.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 4a3dbeb..9c7a57d 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -11,7 +11,7 @@ def __init__(self): def keys(self): keys = [] - module = self.south_module + module_wrapper = self.south_module while module: keys.append(module.key) module = module.north @@ -19,7 +19,7 @@ def keys(self): def modules(self): modules = [] - module = self.south_module + module_wrapper = self.south_module while module: modules.append(module) module = module.north @@ -42,24 +42,24 @@ def add_module( ): north = self.get_module(north_key) south = self.get_module(south_key) - module = ModuleWrapper( + module_wrapper = ModuleWrapper( module, name, north=north, south=south, do_conv=do_conv, ) - self._modules[module.key] = module + self._modules[module_wrapper.key] = module_wrapper if north == self.south_module: - self.south_module = module + self.south_module = module_wrapper if south == self.north_module: - self.north_module = module + self.north_module = module_wrapper - if module.module_type == 'basin': - self.basins.append(module) - elif module.module_type == 'coupler': - self.couplers.append(module) + if module_wrapper.module_type == 'basin': + self.basins.append(module_wrapper) + elif module_wrapper.module_type == 'coupler': + self.couplers.append(module_wrapper) def new_module( self, @@ -80,11 +80,11 @@ def new_module( def get_modules_by_type(self, module_type): modules = [] - module = self.south_module - while module: - if module.module_type == module_type: - modules.append(module) - module = module.north + module_wrapper = self.south_module + while module_wrapper: + if module_wrapper.module_type == module_type: + modules.append(module_wrapper) + module_wrapper = module.north return modules def run( From c6784272c41ae505bc95a1539c4ea9422c7aa55c Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 7 Feb 2020 12:25:53 -0600 Subject: [PATCH 09/28] Add modules to model as full attributes -Allows for easier access to modules -Check that we're not conflicting with existing model attributes --- examples/example_twocol_plusSO_model.py | 74 ++++++++++++++++--------- src/pymoc/model.py | 9 +++ 2 files changed, 57 insertions(+), 26 deletions(-) diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index b76428f..0d2de77 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -124,10 +124,10 @@ def b_north(z): ax1.set_xlabel('$\Psi$', fontsize=14) ax2.set_xlabel('b', fontsize=14) -AMOC = model.get_module('amoc').module -SO = model.get_module('psi_so').module -basin = model.get_module('atlantic_basin').module -north = model.get_module('north_atlantic').module +# AMOC = model.get_module('amoc').module +# model.psi_so = model.get_module('psi_so').module +# basin = model.get_module('atlantic_basin').module +# north = model.get_module('north_atlantic').module for snapshot in model.run( basin_dt=dt, @@ -138,42 +138,59 @@ def b_north(z): ): ax1.cla() ax2.cla() - ax1.plot(AMOC.Psi, AMOC.z, linewidth=0.5, color='r') - ax1.plot(SO.Psi, SO.z, linewidth=0.5, color='m') - ax2.plot(basin.b, basin.z, linewidth=0.5, color='b') - ax2.plot(north.b, north.z, linewidth=0.5, color='c') + ax1.plot(model.amoc.Psi, model.amoc.z, linewidth=0.5, color='r') + ax1.plot(model.psi_so.Psi, model.psi_so.z, linewidth=0.5, color='m') + ax2.plot( + model.atlantic_basin.b, model.atlantic_basin.z, linewidth=0.5, color='b' + ) + ax2.plot( + model.north_atlantic.b, model.north_atlantic.z, linewidth=0.5, color='c' + ) plt.pause(0.01) plt.show() # AMOC = model.get_module('amoc').module -# SO = model.get_module('psi_so').module +# model.psi_so = model.get_module('psi_so').module # basin = model.get_module('atlantic_basin').module # north = model.get_module('north_atlantic').module # # Plot final results over time-iteration plot: -ax1.plot(AMOC.Psi, AMOC.z, linewidth=2, color='r') -ax1.plot(SO.Psi, SO.z, linewidth=2, color='m') -ax1.plot(SO.Psi_Ek, SO.z, linewidth=1, color='m', linestyle='--') -ax1.plot(SO.Psi_GM, SO.z, linewidth=1, color='m', linestyle=':') -ax2.plot(basin.b, basin.z, linewidth=2, color='b') -ax2.plot(north.b, basin.z, linewidth=2, color='c') -ax1.plot(0. * AMOC.z, AMOC.z, linewidth=0.5, color='k') +ax1.plot(model.amoc.Psi, model.amoc.z, linewidth=2, color='r') +ax1.plot(model.psi_so.Psi, model.psi_so.z, linewidth=2, color='m') +ax1.plot( + model.psi_so.Psi_Ek, + model.psi_so.z, + linewidth=1, + color='m', + linestyle='--' +) +ax1.plot( + model.psi_so.Psi_GM, model.psi_so.z, linewidth=1, color='m', linestyle=':' +) +ax2.plot( + model.atlantic_basin.b, model.atlantic_basin.z, linewidth=2, color='b' +) +ax2.plot( + model.north_atlantic.b, model.atlantic_basin.z, linewidth=2, color='c' +) +ax1.plot(0. * model.amoc.z, model.amoc.z, linewidth=0.5, color='k') # #============================================================================== # # Everything below is just to make fancy 2D plots: # # This is mostly an exercise in interpolation, but useful to visualize solutions # # first interpolate buoyancy in channel along constant-slope isopycnals: -bint = Interpolate_channel(y=y, z=z, bs=bs_SO, bn=basin.b) +bint = Interpolate_channel(y=y, z=z, bs=bs_SO, bn=model.atlantic_basin.b) bsouth = bint.gridit() # buoyancy in the basin is all the same: ybasin = np.linspace(200., 12000., 60) -bbasin = np.tile(basin.b, (len(ybasin), 1)) +bbasin = np.tile(model.atlantic_basin.b, (len(ybasin), 1)) # finally, interpolate buoyancy in the north: ynorth = np.linspace(12100., 13000., 10) bnorth = np.zeros((len(ynorth), len(z))) for iz in range(len(z)): bnorth[:, iz] = interp1d([11900., 12000., 12900., 13000.], [ - basin.b[iz], basin.b[iz], north.b[iz], north.b[iz] + model.atlantic_basin.b[iz], model.atlantic_basin.b[iz], + model.north_atlantic.b[iz], model.north_atlantic.b[iz] ], kind='quadratic')(ynorth) # now stick it all together: @@ -184,20 +201,25 @@ def b_north(z): psiarray_iso = np.zeros((len(ynew), len(z))) psiarray_z = np.zeros((len(ynew), len(z))) for iy in range(1, len(y)): - # in the channel, interpolate SO.Psi onto local isopycnal depth: - psiarray_iso[iy, :] = np.interp(bnew[iy, :], basin.b, SO.Psi) + # in the channel, interpolate model.psi_so.Psi onto local isopycnal depth: + psiarray_iso[iy, :] = np.interp( + bnew[iy, :], model.atlantic_basin.b, model.psi_so.Psi + ) psiarray_z[iy, :] = psiarray_iso[iy, :] for iy in range(len(y), len(y) + len(ybasin)): # in the basin, linearly interpolate between Psi_SO and Psi_AMOC: psiarray_iso[iy, :] = ( - ynew[iy] * AMOC.Psibz()[0] + (10000. - ynew[iy]) * SO.Psi + ynew[iy] * model.amoc.Psibz()[0] + (10000. - ynew[iy]) * model.psi_so.Psi + ) / 10000. + psiarray_z[iy, :] = ( + ynew[iy] * model.amoc.Psi + (10000. - ynew[iy]) * model.psi_so.Psi ) / 10000. - psiarray_z[iy, : - ] = (ynew[iy] * AMOC.Psi + (10000. - ynew[iy]) * SO.Psi) / 10000. for iy in range(len(y) + len(ybasin), len(ynew)): # in the north, interpolate AMOC.psib to local isopycnal depth: - psiarray_iso[iy, :] = np.interp(bnew[iy, :], AMOC.bgrid, AMOC.Psib()) - psiarray_z[iy, :] = ((13000. - ynew[iy]) * AMOC.Psi) / 1000. + psiarray_iso[iy, :] = np.interp( + bnew[iy, :], model.amoc.bgrid, model.amoc.Psib() + ) + psiarray_z[iy, :] = ((13000. - ynew[iy]) * model.amoc.Psi) / 1000. psiarray_iso[-1, :] = 0. blevs = np.arange(0.001, 0.03, 0.002) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 9c7a57d..c42d09e 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -49,6 +49,13 @@ def add_module( south=south, do_conv=do_conv, ) + + if hasattr(self, module_wrapper.key): + raise NameError( + 'Cannot use module name ' + name + + ' because it would overwrite an existing key.' + ) + self._modules[module_wrapper.key] = module_wrapper if north == self.south_module: @@ -61,6 +68,8 @@ def add_module( elif module_wrapper.module_type == 'coupler': self.couplers.append(module_wrapper) + setattr(self, module_wrapper.key, module_wrapper.module) + def new_module( self, module_class, From ae359ca90d3d94e2e37b0494a545f8e72429e4a8 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 7 Feb 2020 12:30:46 -0600 Subject: [PATCH 10/28] Allow user to specify do_conv on column initialization --- examples/example_twocol_plusSO_model.py | 7 ++++--- src/pymoc/model.py | 7 +------ src/pymoc/modules/column.py | 8 +++++--- src/pymoc/modules/module_wrapper.py | 4 +--- 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index 0d2de77..5692261 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -102,17 +102,18 @@ def b_north(z): ) # create adv-diff column model instance for basin model.new_module( - Column, { + Column, + { 'z': z, 'kappa': kappa, 'Area': A_north, 'b': b_north, 'bs': bs_north, - 'bbot': bmin + 'bbot': bmin, + 'do_conv': True }, 'North Atlantic', south_key='amoc', - do_conv=True ) fig = plt.figure(figsize=(6, 10)) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index c42d09e..8020d10 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -37,9 +37,7 @@ def get_module(self, key): # return module # return None - def add_module( - self, module, name, north_key=None, south_key=None, do_conv=False - ): + def add_module(self, module, name, north_key=None, south_key=None): north = self.get_module(north_key) south = self.get_module(south_key) module_wrapper = ModuleWrapper( @@ -47,7 +45,6 @@ def add_module( name, north=north, south=south, - do_conv=do_conv, ) if hasattr(self, module_wrapper.key): @@ -77,14 +74,12 @@ def new_module( module_name, north_key=None, south_key=None, - do_conv=False ): self.add_module( module_class(**module_args), module_name, north_key=north_key, south_key=south_key, - do_conv=do_conv ) def get_modules_by_type(self, module_type): diff --git a/src/pymoc/modules/column.py b/src/pymoc/modules/column.py index e98374a..9a71ca2 100644 --- a/src/pymoc/modules/column.py +++ b/src/pymoc/modules/column.py @@ -25,7 +25,8 @@ def __init__( bzbot=None, # bottom strat. as alternative boundary condition (input) b=0.0, # Buoyancy profile (input, output) Area=None, # Horizontal area (can be function of depth) - N2min=1e-7 # Minimum strat. for conv adjustment + N2min=1e-7, # Minimum strat. for conv adjustment + do_conv=False # Whether to do convective adjustment ): r""" Parameters @@ -63,6 +64,7 @@ def __init__( self.bzbot = bzbot self.N2min = N2min + self.do_conv = do_conv self.b = make_array(b, self.z, 'b') @@ -302,7 +304,7 @@ def horadv(self, vdx_in, b_in, dt): self.b[adv_idx] = self.b[adv_idx] + dt * vdx_in[adv_idx] * db[ adv_idx] / self.Area(self.z[adv_idx]) - def timestep(self, wA=0., dt=1., do_conv=False, vdx_in=None, b_in=None): + def timestep(self, wA=0., dt=1., do_conv=None, vdx_in=None, b_in=None): r""" Carry out one timestep integration for the buoyancy profile, accounting for advective, diffusive, and convective effects. @@ -331,6 +333,6 @@ def timestep(self, wA=0., dt=1., do_conv=False, vdx_in=None, b_in=None): self.horadv(vdx_in=vdx_in, b_in=b_in, dt=dt) else: raise TypeError('b_in is needed if vdx_in is provided') - if do_conv: + if do_conv if do_conv is not None else self.do_conv: # do convection: (optional) self.convect() diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index d9617ca..03f827a 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -9,7 +9,6 @@ def __init__( name, north=None, south=None, - do_conv=False, ): if north: north.south = self @@ -19,7 +18,6 @@ def __init__( self.name = name self.north = north self.south = south - self.do_conv = do_conv self.key = name.replace(' ', '_').lower().strip('_') self.do_psi_bz = hasattr(module, 'Psibz') and callable(module.Psibz) self.b_type = 'b' if isinstance(self.module, Column) or isinstance( @@ -39,7 +37,7 @@ def timestep_basin(self, dt=None): wA += self.north.psi[0] * 1e6 if self.south: wA -= self.south.psi[-1] * 1e6 - module.timestep(wA=wA, dt=dt, do_conv=self.do_conv) + module.timestep(wA=wA, dt=dt) def update_coupler(self, dt=None): module = self.module From 75d74ea91c44def36d0b445110b91493bb2052ff Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sat, 8 Feb 2020 09:43:54 -0600 Subject: [PATCH 11/28] Initial pass at generalized neighbors for modules, rather than pure north/south -Still need to check for uniqueness -Still need to enforce that couplers only get two neighbors --- examples/example_twocol_plusSO_model.py | 11 +++--- src/pymoc/__init__.py | 4 +- src/pymoc/model.py | 50 +++++++++++-------------- src/pymoc/modules/__init__.py | 3 +- src/pymoc/modules/module_wrapper.py | 45 +++++++++++----------- src/pymoc/modules/neighbor.py | 10 +++++ src/pymoc/plotting/__init__.py | 2 +- src/pymoc/utils/__init__.py | 2 +- 8 files changed, 63 insertions(+), 64 deletions(-) create mode 100644 src/pymoc/modules/neighbor.py diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index 5692261..4f5e944 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -9,7 +9,7 @@ and Vallis (2012, JPO). ''' from pymoc import model -from pymoc.modules import Psi_Thermwind, Psi_SO, Column +from pymoc.modules import Psi_Thermwind, Psi_SO, Column, Neighbor from pymoc.plotting import Interpolate_channel import numpy as np from matplotlib import pyplot as plt @@ -87,7 +87,7 @@ def b_north(z): 'bbot': bmin }, 'Atlantic Basin', - south_key='psi_so' + neighbors=[Neighbor('psi_so', 'left')] ) # create N.A. overturning model instance model.new_module( @@ -98,12 +98,11 @@ def b_north(z): 'f': 1e-4 }, 'AMOC', - south_key='atlantic_basin' + neighbors=[Neighbor('atlantic_basin', 'left')] ) # create adv-diff column model instance for basin model.new_module( - Column, - { + Column, { 'z': z, 'kappa': kappa, 'Area': A_north, @@ -113,7 +112,7 @@ def b_north(z): 'do_conv': True }, 'North Atlantic', - south_key='amoc', + neighbors=[Neighbor('amoc', 'left')] ) fig = plt.figure(figsize=(6, 10)) diff --git a/src/pymoc/__init__.py b/src/pymoc/__init__.py index d83ad0e..07c260c 100644 --- a/src/pymoc/__init__.py +++ b/src/pymoc/__init__.py @@ -1,5 +1,5 @@ -__version__ = '0.0.1rc5' +__version__ = '0.0.1rc6' import pymoc.modules import pymoc.utils import pymoc.plotting -import pymoc.model \ No newline at end of file +import pymoc.model diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 8020d10..2dee0e2 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -1,10 +1,8 @@ -from pymoc.modules import ModuleWrapper +from pymoc.modules import ModuleWrapper, Neighbor class Model(object): def __init__(self): - self.south_module = None - self.north_module = None self.basins = [] self.couplers = [] self._modules = {} @@ -28,24 +26,25 @@ def modules(self): def get_module(self, key): if not key: return None - return self._modules[key] - # module = self.south_module - # while module and module.key != key: - # module = module.north - # if module and module.key == key: - # return module - # return None - - def add_module(self, module, name, north_key=None, south_key=None): - north = self.get_module(north_key) - south = self.get_module(south_key) - module_wrapper = ModuleWrapper( - module, - name, - north=north, - south=south, - ) + + def add_module(self, module, name, neighbors=[]): + for n in neighbors: + neighbor = self.get_module(n.key) + if not neighbor: + raise KeyError('No module present with key ' + n.key) + n.module_wrapper = neighbor + + module_wrapper = ModuleWrapper(module, name, neighbors) + + for neighbor in neighbors: + neighbor.module_wrapper.neighbors.append( + Neighbor( + module_wrapper.key, + 'left' if n.direction == 'right' else 'right', + module_wrapper=module_wrapper + ) + ) if hasattr(self, module_wrapper.key): raise NameError( @@ -55,11 +54,6 @@ def add_module(self, module, name, north_key=None, south_key=None): self._modules[module_wrapper.key] = module_wrapper - if north == self.south_module: - self.south_module = module_wrapper - if south == self.north_module: - self.north_module = module_wrapper - if module_wrapper.module_type == 'basin': self.basins.append(module_wrapper) elif module_wrapper.module_type == 'coupler': @@ -72,14 +66,12 @@ def new_module( module_class, module_args, module_name, - north_key=None, - south_key=None, + neighbors=[], ): self.add_module( module_class(**module_args), module_name, - north_key=north_key, - south_key=south_key, + neighbors, ) def get_modules_by_type(self, module_type): diff --git a/src/pymoc/modules/__init__.py b/src/pymoc/modules/__init__.py index 0c85360..2048d67 100644 --- a/src/pymoc/modules/__init__.py +++ b/src/pymoc/modules/__init__.py @@ -1,7 +1,8 @@ -__version__ = '0.0.1rc5' +__version__ = '0.0.1rc6' from .column import Column from .equi_column import Equi_Column from .psi_SO import Psi_SO from .psi_thermwind import Psi_Thermwind from .SO_ML import SO_ML +from .neighbor import Neighbor from .module_wrapper import ModuleWrapper diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 03f827a..7c563c3 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -1,30 +1,19 @@ -from datetime import datetime -from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column +import numpy as np +from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column, Neighbor class ModuleWrapper(object): - def __init__( - self, - module, - name, - north=None, - south=None, - ): - if north: - north.south = self - if south: - south.north = self + def __init__(self, module, name, neighbors=[]): self.module = module self.name = name - self.north = north - self.south = south + self.neighbors = neighbors + self.key = name.replace(' ', '_').lower().strip('_') self.do_psi_bz = hasattr(module, 'Psibz') and callable(module.Psibz) self.b_type = 'b' if isinstance(self.module, Column) or isinstance( self.module, Equi_Column ) else 'bs' self.psi = [0, 0] - # self.key = ''.join(['_'+i.lower() if i.isupper() else i for i in name]).lstrip('_').replace(' ', '') @property def module_type(self): @@ -33,20 +22,28 @@ def module_type(self): def timestep_basin(self, dt=None): module = self.module wA = 0.0 - if self.north: - wA += self.north.psi[0] * 1e6 - if self.south: - wA -= self.south.psi[-1] * 1e6 + for neighbor in self.neighbors: + if neighbor.direction == 'right': + wA += neighbor.module_wrapper.psi[0] * 1e6 + else: + wA -= neighbor.module_wrapper.psi[-1] * 1e6 + module.timestep(wA=wA, dt=dt) def update_coupler(self, dt=None): module = self.module - north = self.north - south = self.south + b1 = None + b2 = None + for neighbor in self.neighbors: + if neighbor.direction == 'left': + b1 = neighbor.module_wrapper and neighbor.module_wrapper.b + else: + b2 = neighbor.module_wrapper and neighbor.module_wrapper.b + if self.do_psi_bz: - module.update(b2=north and north.b, b1=south and south.b) + module.update(b1=b1, b2=b2) else: - module.update(b=north and north.b) + module.update(b=b2) module.solve() self.psi = module.Psibz() if self.do_psi_bz else [module.Psi] diff --git a/src/pymoc/modules/neighbor.py b/src/pymoc/modules/neighbor.py new file mode 100644 index 0000000..5044618 --- /dev/null +++ b/src/pymoc/modules/neighbor.py @@ -0,0 +1,10 @@ +class Neighbor(object): + def __init__(self, key, direction, module_wrapper=None): + self.key = key + if direction not in ['left', 'right']: + raise ValueError('Direction must be either left or right.') + self.direction = direction + self.module_wrapper = module_wrapper + + def __eq__(self, other): + return self.key == other.key and self.direction == other.direction diff --git a/src/pymoc/plotting/__init__.py b/src/pymoc/plotting/__init__.py index 0e544ad..4ab50da 100644 --- a/src/pymoc/plotting/__init__.py +++ b/src/pymoc/plotting/__init__.py @@ -1,3 +1,3 @@ -__version__ = '0.0.1rc5' +__version__ = '0.0.1rc6' from .interp_channel import Interpolate_channel from .interp_twocol import Interpolate_twocol diff --git a/src/pymoc/utils/__init__.py b/src/pymoc/utils/__init__.py index 2375d8d..d1600c9 100644 --- a/src/pymoc/utils/__init__.py +++ b/src/pymoc/utils/__init__.py @@ -1,4 +1,4 @@ -__version__ = '0.0.1rc5' +__version__ = '0.0.1rc6' from .check_numpy_version import check_numpy_version from .gridit import gridit from .make_array import make_array From 21ef73b77aae6ee9d55423c1c005ab299909969e Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sat, 8 Feb 2020 10:03:38 -0600 Subject: [PATCH 12/28] Enforce unique neighbor connections --- examples/example_twocol_plusSO_model.py | 1 + src/pymoc/model.py | 11 +++++++++++ 2 files changed, 12 insertions(+) diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index 4f5e944..2f56309 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -100,6 +100,7 @@ def b_north(z): 'AMOC', neighbors=[Neighbor('atlantic_basin', 'left')] ) + # create adv-diff column model instance for basin model.new_module( Column, { diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 2dee0e2..7f696a6 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -1,3 +1,4 @@ +import numpy as np from pymoc.modules import ModuleWrapper, Neighbor @@ -29,6 +30,7 @@ def get_module(self, key): return self._modules[key] def add_module(self, module, name, neighbors=[]): + neighbors = np.unique(neighbors).tolist() for n in neighbors: neighbor = self.get_module(n.key) if not neighbor: @@ -37,6 +39,15 @@ def add_module(self, module, name, neighbors=[]): module_wrapper = ModuleWrapper(module, name, neighbors) + for neighbor in neighbors: + if module_wrapper.key in [ + n.key for n in neighbor.module_wrapper.neighbors + ]: + raise KeyError( + 'Cannot add module ' + module_wrapper.name + ' as a neighbor of ' + + neighbor.module_wrapper.name + ' because they are already coupled.' + ) + for neighbor in neighbors: neighbor.module_wrapper.neighbors.append( Neighbor( From 88fc6d38f337c237cfa52fd28f222fd1892be1b4 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sat, 8 Feb 2020 14:19:15 -0600 Subject: [PATCH 13/28] Configure Neighbors to support two basin problem -We can now support arbitrary (non-mixed layer) basin configurations -Still need to validate uniqueness and overloading of couplers --- examples/twobasin_NadeauJansen.py | 9 +- examples/twobasin_NadeauJansen_model.py | 637 ++++++++++++++++++++++++ src/pymoc/model.py | 7 +- src/pymoc/modules/column.py | 5 +- src/pymoc/modules/module_wrapper.py | 6 +- src/pymoc/modules/psi_thermwind.py | 3 +- 6 files changed, 657 insertions(+), 10 deletions(-) create mode 100644 examples/twobasin_NadeauJansen_model.py diff --git a/examples/twobasin_NadeauJansen.py b/examples/twobasin_NadeauJansen.py index ff12a6c..cdaba00 100644 --- a/examples/twobasin_NadeauJansen.py +++ b/examples/twobasin_NadeauJansen.py @@ -139,8 +139,14 @@ def b_Pac(z): # update buoyancy profile # using isopycnal overturning: wA_Atl = (Psi_iso_Atl + Psi_zonal_Atl - SO_Atl.Psi) * 1e6 + # print('wA_Atl') + # print(wA_Atl) wAN = -Psi_iso_N * 1e6 + # print('wAN') + # print(wAN) wA_Pac = (-Psi_zonal_Pac - SO_Pac.Psi) * 1e6 + # print('wA_Pac') + # print(wA_Pac) Atl.timestep(wA=wA_Atl, dt=dt) north.timestep(wA=wAN, dt=dt, do_conv=True) Pac.timestep(wA=wA_Pac, dt=dt) @@ -169,7 +175,7 @@ def b_Pac(z): ax2.plot(Pac.b, Pac.z, '-g', linewidth=0.5) ax2.plot(north.b, north.z, '-r', linewidth=0.5) plt.pause(0.01) - print "t=", round(ii * dt / 86400 / 360), "years" + print("t=" + str(round(ii * dt / 86400 / 360)) + " years") ax1.plot(AMOC.Psi, AMOC.z, '--r', linewidth=1.5) ax1.plot(ZOC.Psi, ZOC.z, ':c', linewidth=1.5) @@ -180,6 +186,7 @@ def b_Pac(z): ax2.plot(Pac.b, Pac.z, '-g', linewidth=1.5) ax2.plot(north.b, north.z, '-r', linewidth=1.5) +plt.show() # Write out diagnostics (if diag file name is provided): if diag_file is not None: np.savez( diff --git a/examples/twobasin_NadeauJansen_model.py b/examples/twobasin_NadeauJansen_model.py new file mode 100644 index 0000000..d835b37 --- /dev/null +++ b/examples/twobasin_NadeauJansen_model.py @@ -0,0 +1,637 @@ +''' +This script shows an example of a three column model for the +overturning circulation in two basins connected to a channel in the south. +Two columns represent the two basins, while the third column represents +the northern sinking region, connected to one of the basins (teh "Atlantic"). +The meridional overtunning circulation is computed at the northern end of the +first basin (i.e at the interface to the northern sinking region), and at the +the southern end of each basin (i.e. at the interface to the channel). +Moreover, a zonal overturning circulation between the two basins in the channel +is computed (using the thermal wind relation - analog to the computation of the +overturning between the basin and the northern sinking region). +Adiabatic mapping is used for the overturning between the differn columns +''' + +from pymoc import model +from pymoc.modules import Psi_Thermwind, Psi_SO, Column, Neighbor +from pymoc.plotting import Interpolate_channel, Interpolate_twocol +import numpy as np +from matplotlib import pyplot as plt + +diag_file = None + +# boundary conditions: +bs = 0.02 +# the minimum surface buoyancy in the GCM in the North Atlantic is 3.5890e-04 (or 0.1795C) +bs_north = 0.00036 +bAABW = -0.0011 +bbot = min(bAABW, bs_north) + +# S.O. surface boundary conditions and grid: +y = np.asarray(np.linspace(0, 3.e6, 51)) +tau = 0.16 +offset = 0.0345 * (1 - np.cos(np.pi * (5.55e5-1.e5) / 8e6)) +bs_SO = ( + 0.0345 * (1 - np.cos(np.pi * (y-1.e5) / 8e6)) * (y > 5.55e5) + + (bAABW-offset) / 5.55e5 * np.maximum(0, 5.55e5 - y) + offset * + (y < 5.55e5) +) + +# time-stepping parameters: +dt = 86400. * 30. # time-step for vert. adv. diff. calc. +MOC_up_iters = int( + np.floor(2. * 360 * 86400 / dt) +) # multiplier for MOC time-step (MOC is updated every MOC_up_iters time steps) +plot_iters = int( + np.ceil(500 * 360 * 86400 / dt) +) # plotting frequency (in iterations) +total_iters = int( + np.ceil(4000 * 360 * 86400 / dt) +) # total number of timesteps + + +# Effective diffusivity profile +def kappaeff(z): # effective diffusivity profile with tapering in BBL + return 1.0 * ( + 1e-4 * ( + 1.1 - np.tanh( + np.maximum(z + 2000., 0) / 1000. + + np.minimum(z + 2000., 0) / 1300. + ) + ) * (1. - np.maximum(-4000. - z + 600., 0.) / 600.)**2 + ) + + +A_Atl = 7e13 #6e13 +A_north = 5.5e12 +A_Pac = 1.7e14 #1.44e14 + +Lx = 1.3e+07 #(length of the channel) +Latl = 6. / 21. * Lx +Lpac = 15. / 21. * Lx +K = 1800. +N2min = 2e-7 + +# create vertical grid: +z = np.asarray(np.linspace(-4000, 0, 80)) + + +# create initial guess for buoyancy profile in the Atl +def b_Atl(z): + return bs * np.exp(z / 300.) + z / z[0] * bbot + + +#def b_Atl(z): return 0.3*bs+z/z[0]*(bbot-0.3*bs) +def b_Pac(z): + return bs * np.exp(z / 300.) + z / z[0] * bbot + + +m = model.Model() +# create N.A. overturning model instance +m.new_module(Psi_Thermwind, { + 'z': z, +}, 'AMOC') + +# create interbasin zonal overturning model instance +m.new_module(Psi_Thermwind, {'z': z, 'f': 1e-4}, 'ZOC') + +# create S.O. overturning model instance for Atlantic sector +m.new_module( + Psi_SO, { + 'z': z, + 'y': y, + 'b': b_Atl(z), + 'bs': bs_SO, + 'tau': tau, + 'L': Latl, + 'KGM': K + }, 'SO_Atl' +) + +# create S.O. overturning model instance for Pacific sector +m.new_module( + Psi_SO, { + 'z': z, + 'y': y, + 'b': b_Pac(z), + 'bs': bs_SO, + 'tau': tau, + 'L': Lpac, + 'KGM': K + }, 'SO_Pac' +) + +# create adv-diff column model instance for Atl +m.new_module( + Column, { + 'z': z, + 'kappa': kappaeff, + 'b': b_Atl, + 'bs': bs, + 'bbot': bbot, + 'Area': A_Atl, + 'N2min': N2min + }, + 'Atl', + neighbors=[ + Neighbor('amoc', 'right'), + Neighbor('so_atl', 'left'), + Neighbor('zoc', 'right') + ] +) + +# create adv-diff column model instance for northern sinking region +m.new_module( + Column, { + 'z': z, + 'kappa': kappaeff, + 'b': b_Atl, + 'bs': bs_north, + 'bbot': bbot, + 'Area': A_north, + 'N2min': N2min, + 'do_conv': True + }, + 'North', + neighbors=[Neighbor('amoc', 'left')] +) + +# create adv-diff column model instance for Pac +m.new_module( + Column, { + 'z': z, + 'kappa': kappaeff, + 'b': b_Pac, + 'bs': bs, + 'bbot': bbot, + 'Area': A_Pac, + 'N2min': N2min + }, + 'Pac', + neighbors=[Neighbor('zoc', 'left'), + Neighbor('so_pac', 'left')] +) + +# Create figure: +fig = plt.figure(figsize=(6, 9)) +ax1 = fig.add_subplot(111) +ax2 = ax1.twiny() +plt.ylim((-4e3, 0)) +ax1.set_xlim((-20, 30)) +ax2.set_xlim((-0.02, 0.030)) + +for snapshot in m.run( + basin_dt=dt, + coupler_dt=MOC_up_iters, + steps=total_iters, + snapshot_start=plot_iters, + snapshot_interval=plot_iters +): + # update buoyancy profile + # using isopycnal overturning: + # Plot current state: + ax1.plot(m.amoc.Psi, m.amoc.z, '--r', linewidth=0.5) + ax1.plot(m.zoc.Psi, m.zoc.z, ':c', linewidth=0.5) + ax1.plot(m.so_atl.Psi - m.zoc.Psi, m.so_atl.z, '--b', linewidth=0.5) + ax1.plot(m.so_pac.Psi + m.zoc.Psi, m.so_pac.z, '--g', linewidth=0.5) + ax1.plot(m.so_atl.Psi + m.so_pac.Psi, z, '--c', linewidth=0.5) + ax2.plot(m.atl.b, m.atl.z, '-b', linewidth=0.5) + ax2.plot(m.pac.b, m.pac.z, '-g', linewidth=0.5) + ax2.plot(m.north.b, m.north.z, '-r', linewidth=0.5) + plt.pause(0.01) + # print("t=" + str(round(ii * dt / 86400 / 360)) + " years") + +ax1.plot(m.amoc.Psi, m.amoc.z, '--r', linewidth=1.5) +ax1.plot(m.zoc.Psi, m.zoc.z, ':c', linewidth=1.5) +ax1.plot(m.so_atl.Psi - m.zoc.Psi, m.so_atl.z, '--b', linewidth=1.5) +ax1.plot(m.so_pac.Psi + m.zoc.Psi, m.so_pac.z, '--g', linewidth=1.5) +ax1.plot(m.so_atl.Psi + m.so_pac.Psi, z, '--c', linewidth=1.5) +ax2.plot(m.atl.b, m.atl.z, '-b', linewidth=1.5) +ax2.plot(m.pac.b, m.pac.z, '-g', linewidth=1.5) +ax2.plot(m.north.b, m.north.z, '-r', linewidth=1.5) + +plt.show() +# Write out diagnostics (if diag file name is provided): +# if diag_file is not None: +# np.savez( +# diag_file, +# b_Atl=Atl.b, +# b_Pac=Pac.b, +# b_north=north.b, +# Psi_SO_Atl=SO_Atl.Psi, +# Psi_SO_Pac=SO_Pac.Psi, +# Psi_ZOC=ZOC.Psi, +# Psib_ZOC=ZOC.Psib(), +# bgrid_ZOC=ZOC.bgrid, +# Psi_AMOC=AMOC.Psi, +# Psib_AMOC=AMOC.Psib(), +# bgrid_AMOC=AMOC.bgrid +# ) + + # #***************************************************************************** + # # Below is all just for fancy plots + # #***************************************************************************** + + # blevs = np.array([-0.004, -0.002, 0.0, 0.004, 0.01, 0.018]) + # plevs = np.arange(-21., 23., 2.0) + # nb = 500 + + # b_basin = (A_Atl * Atl.b + A_Pac * Pac.b) / (A_Atl+A_Pac) + # b_north = north.b + + # # Compute total SO residual overturning by first interpolating circulation + # # in both sectors to mean buoyancy profile values and then summing (at constant b). + # PsiSO = ( + # np.interp(b_basin, Atl.b, SO_Atl.Psi) + + # np.interp(b_basin, Pac.b, SO_Pac.Psi) + # ) + # AMOC_bbasin = np.interp(b_basin, AMOC.bgrid, AMOC.Psib(nb=nb)) + # # this is to Plot adiabatic overturning only: + # PsiSO_bgrid = ( + # np.interp(AMOC.bgrid, Atl.b, SO_Atl.Psi) + + # np.interp(AMOC.bgrid, Pac.b, SO_Pac.Psi) + # ) + + # l = y[-1] + + # bs = 1. * bs_SO + # bn = 1. * b_basin + # bs[-1] = 1. * bn[-1] + # # if bs is tiny bit warmer there is a zero slope point at the surface which makes the interpolation fail... + + # # first interpolate buoyancy in channel along constant-slope isopycnals: + # bint = Interpolate_channel(y=y, z=z, bs=bs, bn=bn) + # bsouth = bint.gridit() + # # buoyancy in the basin is all the same: + # lbasin = 11000. + # ltrans = 1500. + # lnorth = 400. + # lchannel = l / 1e3 + # ybasin = np.linspace(0, lbasin, 60) + lchannel + # bbasin = np.tile(b_basin, (len(ybasin), 1)) + # bAtl = np.tile(Atl.b, (len(ybasin), 1)) + # bPac = np.tile(Pac.b, (len(ybasin), 1)) + # # interpolate buoyancy in northern ransition region: + # ytrans = np.linspace(ltrans / 20., ltrans, 20) + lchannel + lbasin + # bn = b_north.copy() + # bn[0] = b_basin[0] + # # Notice that the interpolation procedure assumes that the bottom + # #buoyancies in both colums match - which may not be exactly the case depending + # # on when in teh time-step data is saved + # bint = Interpolate_twocol( + # y=ytrans*1000. - ytrans[0] * 1000., z=z, bs=Atl.b, bn=bn + # ) + # btrans = bint.gridit() + # # finally set buyancy in northern deep water formation region: + # ynorth = np.linspace(lnorth / 20., lnorth, 20) + lchannel + lbasin + ltrans + # bnorth = np.tile(bn, (len(ynorth), 1)) + # # now stick it all together: + # ynew = np.concatenate((y / 1e3, ybasin, ytrans, ynorth)) + # bnew = np.concatenate((bsouth, bbasin, btrans, bnorth)) + # bnew_Atl = np.concatenate((bsouth, bAtl, btrans, bnorth)) + # bnew_Pac = np.concatenate((bsouth, bPac, np.nan * btrans, np.nan * bnorth)) + + # # Compute z-coordinate and residual overturning streamfunction at all latitudes: + # psiarray_b = np.zeros((len(ynew), len(z))) # overturning in b-coordinates + # psiarray_b_Atl = np.zeros((len(ynew), len(z)) + # ) # overturning in b-coordinates + # psiarray_b_Pac = np.zeros((len(ynew), len(z)) + # ) # overturning in b-coordinates + # psiarray_Atl = np.zeros((len(ynew), len(z)) + # ) # "residual" overturning in ATl. + # psiarray_Pac = np.zeros((len(ynew), len(z)) + # ) # "residual" overturning in Pac. + # psiarray_z = np.zeros((len(ynew), len(z))) # z-space, "eulerian" overturning + # psiarray_z_Atl = np.zeros((len(ynew), len(z)) + # ) # z-space, "eulerian" overturning + # psiarray_z_Pac = np.zeros((len(ynew), len(z)) + # ) # z-space, "eulerian" overturning + # for iy in range(1, len(y)): + # # in the channel, interpolate PsiSO onto local isopycnal depth: + # psiarray_z[iy, :] = np.interp(bnew[iy, :], b_basin, SO_Atl.Psi + SO_Pac.Psi) + # psiarray_z_Atl[iy, :] = psiarray_z[iy, :] + # psiarray_z_Pac[iy, :] = psiarray_z[iy, :] + # psiarray_b[iy, b_basin < bs_SO[iy]] = PsiSO[b_basin < bs_SO[iy]] + # psiarray_Atl[iy, :] = np.interp(bnew[iy, :], b_basin, PsiSO) + # psiarray_b_Atl[iy, :] = psiarray_b[iy, :] + # psiarray_Pac[iy, :] = np.interp(bnew[iy, :], b_basin, PsiSO) + # psiarray_b_Pac[iy, :] = psiarray_b[iy, :] + # for iy in range(len(y), len(y) + len(ybasin)): + # # in the basin, linearly interpolate between Psi_SO and Psi_AMOC: + # psiarray_z[iy, :] = ((ynew[iy] - lchannel) * AMOC.Psi + + # (lchannel + lbasin - ynew[iy]) * + # (SO_Atl.Psi + SO_Pac.Psi)) / lbasin + # psiarray_z_Atl[iy, :] = ((ynew[iy] - lchannel) * AMOC.Psi + + # (lchannel + lbasin - ynew[iy]) * + # (SO_Atl.Psi - ZOC.Psi)) / lbasin + # psiarray_z_Pac[iy, :] = ((lchannel + lbasin - ynew[iy]) * + # (SO_Pac.Psi + ZOC.Psi)) / lbasin + # psiarray_b[iy, :] = ((ynew[iy] - lchannel) * AMOC_bbasin + + # (lchannel + lbasin - ynew[iy]) * PsiSO) / lbasin + # psiarray_Atl[iy, :] = ((ynew[iy] - lchannel) * AMOC.Psibz(nb=nb)[0] + + # (lchannel + lbasin - ynew[iy]) * + # (SO_Atl.Psi - ZOC.Psibz()[0])) / lbasin + # psiarray_b_Atl[iy, :] = ((ynew[iy] - lchannel) * AMOC_bbasin + + # (lchannel + lbasin - ynew[iy]) * ( + # np.interp(b_basin, Atl.b, SO_Atl.Psi) - + # np.interp(b_basin, ZOC.bgrid, ZOC.Psib()) + # )) / lbasin + # psiarray_Pac[iy, :] = ((lchannel + lbasin - ynew[iy]) * + # (SO_Pac.Psi + ZOC.Psibz()[1])) / lbasin + # psiarray_b_Pac[iy, :] = ((lchannel + lbasin - ynew[iy]) * ( + # np.interp(b_basin, Pac.b, SO_Pac.Psi) + + # np.interp(b_basin, ZOC.bgrid, ZOC.Psib()) + # )) / lbasin + # for iy in range(len(y) + len(ybasin), len(y) + len(ybasin) + len(ytrans)): + # # in the northern transition region keep psi_z constant + # # and psi_b constant on non-outcropped isopycnals (but zero on outcropped + # # isopycnals): + # psiarray_z[iy, :] = AMOC.Psi + # psiarray_z_Atl[iy, :] = AMOC.Psi + # psiarray_z_Pac[iy, :] = np.NaN + # psiarray_b[iy, b_basin < bnew[iy, -1]] = AMOC_bbasin[b_basin < bnew[iy, -1]] + # psiarray_Atl[iy, :] = np.interp(bnew[iy, :], AMOC.bgrid, AMOC.Psib(nb=nb)) + # psiarray_b_Atl[iy, b_basin < bnew[iy, -1] + # ] = psiarray_b[iy, b_basin < bnew[iy, -1]] + # psiarray_Pac[iy, :] = np.NaN + # psiarray_b_Pac[iy, :] = np.NaN + # for iy in range(len(y) + len(ybasin) + len(ytrans), len(ynew)): + # # in the northern sinking region, all psi decrease linearly to zero: + # psiarray_z[iy, :] = ((lchannel + lbasin + ltrans + lnorth - ynew[iy]) * + # AMOC.Psi) / lnorth + # psiarray_z_Atl[iy, :] = ((lchannel + lbasin + ltrans + lnorth - ynew[iy]) * + # AMOC.Psi) / lnorth + # psiarray_z_Pac[iy, :] = np.NaN + # psiarray_b[iy, b_basin < bnew[iy, -1]] = ( + # (lchannel + lbasin + ltrans + lnorth - ynew[iy]) * + # AMOC_bbasin[b_basin < bnew[iy, -1]] + # ) / lnorth + # psiarray_Atl[iy, :] = ((lchannel + lbasin + ltrans + lnorth - ynew[iy]) * + # AMOC.Psibz(nb=nb)[1]) / lnorth + # psiarray_b_Atl[iy, b_basin < bnew[iy, -1] + # ] = psiarray_b[iy, b_basin < bnew[iy, -1]] + # psiarray_Pac[iy, :] = np.NaN + # psiarray_b_Pac[iy, :] = np.NaN + + # # plot z-coordinate overturning and buoyancy structure: + + # fig = plt.figure(figsize=(10.8, 6.8)) + # ax1 = fig.add_axes([0.1, .57, .33, .36]) + # ax2 = ax1.twiny() + # plt.ylim((-4e3, 0)) + # ax1.set_ylabel('Depth [m]', fontsize=13) + # ax1.set_xlim((-20, 30)) + # ax2.set_xlim((-0.02, 0.030)) + # ax1.set_xlabel('$\Psi$ [SV]', fontsize=13) + # ax2.set_xlabel('$b_B$ [m s$^{-2}$]', fontsize=13) + # ax1.plot(AMOC.Psi, AMOC.z, '--r', linewidth=1.5) + # ax1.plot(ZOC.Psi, ZOC.z, ':c', linewidth=1.5) + # ax1.plot(SO_Atl.Psi - ZOC.Psi, SO_Atl.z, '--b', linewidth=1.5) + # ax1.plot(SO_Pac.Psi + ZOC.Psi, SO_Pac.z, '--g', linewidth=1.5) + # ax1.plot(SO_Atl.Psi + SO_Pac.Psi, z, '--k', linewidth=1.5) + # ax2.plot(Atl.b, Atl.z, '-b', linewidth=1.5) + # ax2.plot(Pac.b, Pac.z, '-g', linewidth=1.5) + # ax2.plot(north.b, north.z, '-r', linewidth=1.5) + # h1, l1 = ax1.get_legend_handles_labels() + # h2, l2 = ax2.get_legend_handles_labels() + # ax1.legend(h1 + h2, l1 + l2, loc=4, frameon=False) + # ax1.plot(0. * z, z, linewidth=0.5, color='k', linestyle=':') + + # ax1 = fig.add_subplot(222) + # CS = ax1.contour( + # ynew, + # z, + # bnew.transpose() / 2e-3, + # levels=blevs / 2e-3, + # colors='k', + # linewidths=1.0, + # linestyles='solid' + # ) + # ax1.clabel(CS, fontsize=10, fmt='%1.0f') + # ax1.contour( + # ynew, z, psiarray_z.transpose(), levels=plevs, colors='k', linewidths=0.5 + # ) + # CS = ax1.contourf( + # ynew, + # z, + # psiarray_z.transpose(), + # levels=plevs, + # cmap=plt.cm.bwr, + # vmin=-20, + # vmax=20 + # ) + # ax1.set_xlim([0, ynew[-1]]) + # ax1.set_xlabel('y [km]', fontsize=12) + # ax1.set_ylabel('Depth [m]', fontsize=12) + # ax1.set_title('Global', fontsize=12) + + # ax1 = fig.add_subplot(223) + # CS = ax1.contour( + # ynew, + # z, + # bnew_Pac.transpose() / 2e-3, + # levels=blevs / 2e-3, + # colors='k', + # linewidths=1.0, + # linestyles='solid' + # ) + # ax1.clabel(CS, fontsize=10, fmt='%1.0f') + # ax1.contour( + # ynew, + # z, + # psiarray_z_Pac.transpose(), + # levels=plevs, + # colors='k', + # linewidths=0.5 + # ) + # CS = ax1.contourf( + # ynew, + # z, + # psiarray_z_Pac.transpose(), + # levels=plevs, + # cmap=plt.cm.bwr, + # vmin=-20, + # vmax=20 + # ) + # ax1.set_xlim([0, ynew[-1]]) + # ax1.set_xlabel('y [km]', fontsize=12) + # ax1.set_title('Pacific', fontsize=12) + # plt.gca().invert_xaxis() + + # ax1 = fig.add_subplot(224) + # CS = ax1.contour( + # ynew, + # z, + # bnew_Atl.transpose() / 2e-3, + # levels=blevs / 2e-3, + # colors='k', + # linewidths=1.0, + # linestyles='solid' + # ) + # ax1.clabel(CS, fontsize=10, fmt='%1.0f') + # ax1.contour( + # ynew, + # z, + # psiarray_z_Atl.transpose(), + # levels=plevs, + # colors='k', + # linewidths=0.5 + # ) + # CS = ax1.contourf( + # ynew, + # z, + # psiarray_z_Atl.transpose(), + # levels=plevs, + # cmap=plt.cm.bwr, + # vmin=-20, + # vmax=20 + # ) + # ax1.set_xlim([0, ynew[-1]]) + # ax1.set_xlabel('y [km]', fontsize=12) + # ax1.set_ylabel('Depth [m]', fontsize=12) + # ax1.set_title('Atlantic', fontsize=12) + + # cbar_ax = fig.add_axes([0.94, 0.15, 0.015, 0.7]) + # fig.tight_layout() + # fig.subplots_adjust(right=0.92) + # fig.colorbar( + # CS, cax=cbar_ax, ticks=np.arange(-20, 21, 5), orientation='vertical' + # ) + + # # Plot Isopycnal overturning + + # fig = plt.figure(figsize=(10.8, 6.8)) + # ax1 = fig.add_axes([0.1, .57, .33, .36]) + # ax2 = ax1.twiny() + # plt.ylim((-4e3, 0)) + # ax1.set_ylabel('b [m s$^{-2}$]', fontsize=13) + # ax1.set_xlim((-20, 30)) + # ax2.set_xlim((-0.02, 0.030)) + # ax1.set_xlabel('$\Psi$ [SV]', fontsize=13) + # ax2.set_xlabel('$b_B$ [m s$^{-2}$]', fontsize=13) + # ax1.plot(AMOC_bbasin, z, '--r', linewidth=1.5) + # ax1.plot(np.interp(b_basin, ZOC.bgrid, ZOC.Psib()), z, ':c', linewidth=1.5) + # ax1.plot( + # np.interp(b_basin, Atl.b, SO_Atl.Psi) - + # np.interp(b_basin, ZOC.bgrid, ZOC.Psib()), + # z, + # '--b', + # linewidth=1.5 + # ) + # ax1.plot( + # np.interp(b_basin, Pac.b, SO_Pac.Psi) + + # np.interp(b_basin, ZOC.bgrid, ZOC.Psib()), + # z, + # '--g', + # linewidth=1.5 + # ) + # ax1.plot(PsiSO, z, '--k', linewidth=1.5) + # ax2.plot(Atl.b, Atl.z, '-b', linewidth=1.5) + # ax2.plot(Pac.b, Pac.z, '-g', linewidth=1.5) + # ax2.plot(north.b, north.z, '-r', linewidth=1.5) + # h1, l1 = ax1.get_legend_handles_labels() + # h2, l2 = ax2.get_legend_handles_labels() + # ax1.legend(h1 + h2, l1 + l2, loc=4, frameon=False) + # ax1.plot(0. * b_basin, b_basin, linewidth=0.5, color='k', linestyle=':') + # ax1.set_yticks( + # np.interp([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001], + # b_basin, z) + # ) + # ax1.set_yticklabels([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001]) + + # ax1 = fig.add_subplot(222) + # ax1.contour( + # ynew, z, psiarray_b.transpose(), levels=plevs, colors='k', linewidths=0.5 + # ) + # CS = ax1.contourf( + # ynew, + # z, + # psiarray_b.transpose(), + # levels=plevs, + # cmap=plt.cm.bwr, + # vmin=-20, + # vmax=20 + # ) + # ax1.set_xlim([0, ynew[-1]]) + # ax1.set_xlabel('y [km]', fontsize=12) + # ax1.set_ylabel('b [m s$^{-2}$]', fontsize=12) + # ax1.set_title('Global', fontsize=12) + # ax1.set_yticks( + # np.interp([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001], + # b_basin, z) + # ) + # ax1.set_yticklabels([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001]) + + # ax1 = fig.add_subplot(223) + # ax1.contour( + # ynew, + # z, + # psiarray_b_Pac.transpose(), + # levels=plevs, + # colors='k', + # linewidths=0.5 + # ) + # CS = ax1.contourf( + # ynew, + # z, + # psiarray_b_Pac.transpose(), + # levels=plevs, + # cmap=plt.cm.bwr, + # vmin=-20, + # vmax=20 + # ) + # ax1.plot( + # np.array([y[-1] / 1000., y[-1] / 1000.]), + # np.array([-4000., 0.]), + # color='k' + # ) + # ax1.set_xlim([0, ynew[-1]]) + # ax1.set_xlabel('y [km]', fontsize=12) + # ax1.set_ylabel('b [m s$^{-2}$]', fontsize=12) + # ax1.set_title('Pacific', fontsize=12) + # ax1.set_yticks( + # np.interp([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001], + # b_basin, z) + # ) + # ax1.set_yticklabels([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001]) + # plt.gca().invert_xaxis() + + # ax1 = fig.add_subplot(224) + # ax1.contour( + # ynew, + # z, + # psiarray_b_Atl.transpose(), + # levels=plevs, + # colors='k', + # linewidths=0.5 + # ) + # CS = ax1.contourf( + # ynew, + # z, + # psiarray_b_Atl.transpose(), + # levels=plevs, + # cmap=plt.cm.bwr, + # vmin=-20, + # vmax=20 + # ) + # ax1.plot( + # np.array([y[-1] / 1000., y[-1] / 1000.]), + # np.array([-4000., 0.]), + # color='k' + # ) + # ax1.set_xlim([0, ynew[-1]]) + # ax1.set_xlabel('y [km]', fontsize=12) + # ax1.set_ylabel('b [m s$^{-2}$]', fontsize=12) + # ax1.set_title('Atlantic', fontsize=12) + # ax1.set_yticks( + # np.interp([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001], + # b_basin, z) + # ) + # ax1.set_yticklabels([0.02, 0.005, 0.002, 0.001, 0.0005, 0., -0.0005, -0.001]) + + # cbar_ax = fig.add_axes([0.94, 0.15, 0.015, 0.7]) + # fig.tight_layout() + # fig.subplots_adjust(right=0.92) + # fig.colorbar( + # CS, cax=cbar_ax, ticks=np.arange(-20, 21, 5), orientation='vertical' + # ) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 7f696a6..d438bc0 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -30,7 +30,7 @@ def get_module(self, key): return self._modules[key] def add_module(self, module, name, neighbors=[]): - neighbors = np.unique(neighbors).tolist() + # neighbors = np.unique(neighbors).tolist() for n in neighbors: neighbor = self.get_module(n.key) if not neighbor: @@ -52,7 +52,7 @@ def add_module(self, module, name, neighbors=[]): neighbor.module_wrapper.neighbors.append( Neighbor( module_wrapper.key, - 'left' if n.direction == 'right' else 'right', + 'left' if neighbor.direction == 'right' else 'right', module_wrapper=module_wrapper ) ) @@ -102,6 +102,9 @@ def run( snapshot_start=None, snapshot_interval=None ): + for coupler in self.couplers: + coupler.update_coupler() + for i in range(0, steps): self.timestep(i, basin_dt, coupler_dt=coupler_dt) if self.snapshot(i, snapshot_start, snapshot_interval): diff --git a/src/pymoc/modules/column.py b/src/pymoc/modules/column.py index 4408189..6ee57b1 100644 --- a/src/pymoc/modules/column.py +++ b/src/pymoc/modules/column.py @@ -232,7 +232,7 @@ def vertadvdiff(self, wA, dt, do_conv=False): # apply boundary conditions: # TODO: Make this check the module level do_conv as well - if not do_conv: # if we use convection, upper BC is already applied there + if not self.do_conv: # if we use convection, upper BC is already applied there self.b[-1] = self.bs self.b[0] = ( self.bbot if self.bzbot is None else self.b[1] - self.bzbot * dz[0] @@ -341,7 +341,8 @@ def timestep(self, wA=0., dt=1., do_conv=None, vdx_in=None, b_in=None): """ #TODO: Remove backwards compatibility for method level do_conv from module and examples. Adds unnecesarry complexity. - if (do_conv is None and self.do_conv) or do_conv: + if self.do_conv: + # if (do_conv is None and self.do_conv) or do_conv: # do convection: (optional) self.convect() diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 7c563c3..a7cb4a0 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -3,10 +3,10 @@ class ModuleWrapper(object): - def __init__(self, module, name, neighbors=[]): + def __init__(self, module, name, neighbors=None): self.module = module self.name = name - self.neighbors = neighbors + self.neighbors = neighbors or [] self.key = name.replace(' ', '_').lower().strip('_') self.do_psi_bz = hasattr(module, 'Psibz') and callable(module.Psibz) @@ -30,7 +30,7 @@ def timestep_basin(self, dt=None): module.timestep(wA=wA, dt=dt) - def update_coupler(self, dt=None): + def update_coupler(self): module = self.module b1 = None b2 = None diff --git a/src/pymoc/modules/psi_thermwind.py b/src/pymoc/modules/psi_thermwind.py index 2047a91..b378281 100644 --- a/src/pymoc/modules/psi_thermwind.py +++ b/src/pymoc/modules/psi_thermwind.py @@ -32,7 +32,7 @@ def __init__( f=1.2e-4, # Coriolis parameter (input) z=None, # grid (input) sol_init=None, # Initial conditions for ODE solver (input) - b1=None, # Buoyancy in the basin (input, output) + b1=0., # Buoyancy in the basin (input, output) b2=0., # Buoyancy in the deep water formation region (input, output) ): r""" @@ -62,7 +62,6 @@ def __init__( self.b1 = make_func(b1, self.z, 'b1') self.b2 = make_func(b2, self.z, 'b2') - # Set initial conditions for BVP solver if sol_init is None: self.sol_init = np.zeros((2, nz)) From b218338ead757d295107f0d413a572d7c776205e Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sat, 8 Feb 2020 14:49:10 -0600 Subject: [PATCH 14/28] Enforce coupler connection rules -Coupler can be connected to no more than two modules -Coupler cannot be connected in the same direction multiple times --- src/pymoc/model.py | 40 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index d438bc0..d24ac73 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -29,8 +29,14 @@ def get_module(self, key): return None return self._modules[key] - def add_module(self, module, name, neighbors=[]): - # neighbors = np.unique(neighbors).tolist() + def add_module(self, module, name, neighbors=None): + neighbors = neighbors or [] + if len(np.unique([n.key + for n in neighbors])) < len([n.key for n in neighbors]): + raise ValueError( + 'Cannot link basins multiple times. Please check your configuration.' + ) + for n in neighbors: neighbor = self.get_module(n.key) if not neighbor: @@ -38,6 +44,16 @@ def add_module(self, module, name, neighbors=[]): n.module_wrapper = neighbor module_wrapper = ModuleWrapper(module, name, neighbors) + if module_wrapper.module_type == 'coupler': + if len(neighbors) > 2: + raise ValueError( + 'Streamfunctions cannot connect more than two basins. Please check your configuration.' + ) + if sum(n.direction == 'left' for n in neighbors + ) > 1 or sum(n.direction == 'right' for n in neighbors) > 1: + raise ValueError( + 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' + ) for neighbor in neighbors: if module_wrapper.key in [ @@ -56,6 +72,19 @@ def add_module(self, module, name, neighbors=[]): module_wrapper=module_wrapper ) ) + if neighbor.module_wrapper.module_type == 'coupler': + if len(neighbor.module_wrapper.neighbors) > 2: + raise ValueError( + 'Streamfunctions cannot connect more than two basins. Please check your configuration.' + ) + if sum( + n.direction == 'left' for n in neighbor.module_wrapper.neighbors + ) > 1 or sum( + n.direction == 'right' for n in neighbor.module_wrapper.neighbors + ) > 1: + raise ValueError( + 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' + ) if hasattr(self, module_wrapper.key): raise NameError( @@ -63,6 +92,13 @@ def add_module(self, module, name, neighbors=[]): ' because it would overwrite an existing key.' ) + if module_wrapper.module_type == 'coupler' and len( + module_wrapper.neighbors + ) > 2: + raise ValueError( + 'Streamfunctions cannot connect more than two basins. Please check your configuration.' + ) + self._modules[module_wrapper.key] = module_wrapper if module_wrapper.module_type == 'basin': From 15a1d582818e81097b57515e3970cb92d60520a8 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sun, 9 Feb 2020 12:12:59 -0600 Subject: [PATCH 15/28] Remove superfluous checks --- src/pymoc/model.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index d24ac73..d0e9db9 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -45,10 +45,7 @@ def add_module(self, module, name, neighbors=None): module_wrapper = ModuleWrapper(module, name, neighbors) if module_wrapper.module_type == 'coupler': - if len(neighbors) > 2: - raise ValueError( - 'Streamfunctions cannot connect more than two basins. Please check your configuration.' - ) + # We don't need to explicitly check that a coupler has two or fewer neighbors, as the enforcement of only one left, only one right, and only being able to point left or right implicitly enforces that condition. if sum(n.direction == 'left' for n in neighbors ) > 1 or sum(n.direction == 'right' for n in neighbors) > 1: raise ValueError( @@ -73,10 +70,6 @@ def add_module(self, module, name, neighbors=None): ) ) if neighbor.module_wrapper.module_type == 'coupler': - if len(neighbor.module_wrapper.neighbors) > 2: - raise ValueError( - 'Streamfunctions cannot connect more than two basins. Please check your configuration.' - ) if sum( n.direction == 'left' for n in neighbor.module_wrapper.neighbors ) > 1 or sum( From d1858d0c04970defe7c7f49551555ceb756630e2 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sun, 9 Feb 2020 14:56:12 -0600 Subject: [PATCH 16/28] Clean up neighbor based model validation --- src/pymoc/model.py | 81 ++++++----------------------- src/pymoc/modules/module_wrapper.py | 35 +++++++++++++ 2 files changed, 51 insertions(+), 65 deletions(-) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index d0e9db9..2a03174 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -9,30 +9,17 @@ def __init__(self): self._modules = {} def keys(self): - keys = [] - module_wrapper = self.south_module - while module: - keys.append(module.key) - module = module.north - return keys - - def modules(self): - modules = [] - module_wrapper = self.south_module - while module: - modules.append(module) - module = module.north - return modules + return self._modules.keys() def get_module(self, key): if not key: return None return self._modules[key] - def add_module(self, module, name, neighbors=None): - neighbors = neighbors or [] - if len(np.unique([n.key - for n in neighbors])) < len([n.key for n in neighbors]): + def validate_neighbors_input(self, neighbors): + neighbor_keys = [n.key for n in neighbors] + distinct_neighbor_keys = np.unique(neighbor_keys) + if len(neighbor_keys) > len(distinct_neighbor_keys): raise ValueError( 'Cannot link basins multiple times. Please check your configuration.' ) @@ -40,60 +27,26 @@ def add_module(self, module, name, neighbors=None): for n in neighbors: neighbor = self.get_module(n.key) if not neighbor: - raise KeyError('No module present with key ' + n.key) - n.module_wrapper = neighbor - - module_wrapper = ModuleWrapper(module, name, neighbors) - if module_wrapper.module_type == 'coupler': - # We don't need to explicitly check that a coupler has two or fewer neighbors, as the enforcement of only one left, only one right, and only being able to point left or right implicitly enforces that condition. - if sum(n.direction == 'left' for n in neighbors - ) > 1 or sum(n.direction == 'right' for n in neighbors) > 1: - raise ValueError( - 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' - ) - - for neighbor in neighbors: - if module_wrapper.key in [ - n.key for n in neighbor.module_wrapper.neighbors - ]: raise KeyError( - 'Cannot add module ' + module_wrapper.name + ' as a neighbor of ' + - neighbor.module_wrapper.name + ' because they are already coupled.' + 'No module present with key ' + n.key + + ', cannot set as neighbor of ' + name ) + n.module_wrapper = neighbor - for neighbor in neighbors: - neighbor.module_wrapper.neighbors.append( - Neighbor( - module_wrapper.key, - 'left' if neighbor.direction == 'right' else 'right', - module_wrapper=module_wrapper - ) - ) - if neighbor.module_wrapper.module_type == 'coupler': - if sum( - n.direction == 'left' for n in neighbor.module_wrapper.neighbors - ) > 1 or sum( - n.direction == 'right' for n in neighbor.module_wrapper.neighbors - ) > 1: - raise ValueError( - 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' - ) + def add_module(self, module, name, neighbors=None): + neighbors = neighbors or [] + self.validate_neighbors_input(neighbors) + module_wrapper = ModuleWrapper(module, name, neighbors) + module_wrapper.validate_neighbor_links() + module_wrapper.backlink_neighbors() if hasattr(self, module_wrapper.key): raise NameError( 'Cannot use module name ' + name + - ' because it would overwrite an existing key.' - ) - - if module_wrapper.module_type == 'coupler' and len( - module_wrapper.neighbors - ) > 2: - raise ValueError( - 'Streamfunctions cannot connect more than two basins. Please check your configuration.' + ' because it would overwrite an existing key or model property.' ) self._modules[module_wrapper.key] = module_wrapper - if module_wrapper.module_type == 'basin': self.basins.append(module_wrapper) elif module_wrapper.module_type == 'coupler': @@ -116,11 +69,9 @@ def new_module( def get_modules_by_type(self, module_type): modules = [] - module_wrapper = self.south_module - while module_wrapper: + for module_wrapper in self._modules.values(): if module_wrapper.module_type == module_type: modules.append(module_wrapper) - module_wrapper = module.north return modules def run( diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index a7cb4a0..5ef3e4b 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -50,3 +50,38 @@ def update_coupler(self): @property def b(self): return getattr(self.module, self.b_type) + + def validate_neighbors(self, backlink=False): + if not backlink: + for neighbor in self.neighbors: + linked_neighbor_keys = [ + n.key for n in neighbor.module_wrapper.neighbors + ] + if self.key in linked_neighbor_keys: + raise KeyError( + 'Cannot add module ' + self.name + ' as a neighbor of ' + + neighbor.module_wrapper.name + + ' because they are already coupled.' + ) + + if self.module_type == 'coupler': + # We don't need to explicitly check that a coupler has two or fewer neighbors, as the enforcement of only one left, only one right, and only being able to point left or right implicitly enforces that condition. + left_neighbor_count = sum(n.direction == 'left' for n in self.neighbors) + right_neighbor_count = sum( + n.direction == 'right' for n in self.neighbors + ) + if left_neighbor_count > 1 or right_neighbor_count > 1: + raise ValueError( + 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' + ) + + def backlink_neighbor_links(self): + for neighbor in self.neighbors: + neighbor.module_wrapper.neighbors.append( + Neighbor( + self.key, + 'left' if neighbor.direction == 'right' else 'right', + module_wrapper=self + ) + ) + neighbor.module_wrapper.validate_neighbors(backlink=True) From 0423ccb9ba2396f7fd0a784896dcd0bc3fd16ecf Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sun, 9 Feb 2020 14:57:58 -0600 Subject: [PATCH 17/28] remove unused method --- src/pymoc/model.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 2a03174..c26d873 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -67,13 +67,6 @@ def new_module( neighbors, ) - def get_modules_by_type(self, module_type): - modules = [] - for module_wrapper in self._modules.values(): - if module_wrapper.module_type == module_type: - modules.append(module_wrapper) - return modules - def run( self, steps, From dd5f65fda5db8d7709989e520c1fdaeaf47e8f9b Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Tue, 25 Feb 2020 06:59:02 -0500 Subject: [PATCH 18/28] Move thermalwind documentation to right/left convention --- docs/model-physics/thermal-wind-closure.rst | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/docs/model-physics/thermal-wind-closure.rst b/docs/model-physics/thermal-wind-closure.rst index b2fa19b..48e0356 100644 --- a/docs/model-physics/thermal-wind-closure.rst +++ b/docs/model-physics/thermal-wind-closure.rst @@ -6,11 +6,12 @@ North Atlantic deep water formation region, is represented by a Thermal wind closure .. math:: - \partial_{zz}\Psi^z_N\left(z\right)=\frac{1}{f}\left[b_B\left(z\right)-b_N\left(z\right)\right] + \partial_{zz}\Psi\left(z\right)=\frac{1}{f}\left[b_{left}\left(z\right)-b_{right}\left(z\right)\right] -Where :math:`b_B(z)` is the basin interior buoyancy, :math:`b_N(z)` is the northern -basin buoyancy, :math:`f` is the coriolis parameter, and :math:`\Psi^z_N(z)` is -the **zonal** overturning streamfunction in the northern basin. +Where :math:`b_{left}(z)` is the lefthand (basin interior) buoyancy, :math:`b_{right}(z)` is the righthand (northern +basin) buoyancy, :math:`f` is the coriolis parameter, and :math:`\Psi(z)` is +the **zonal** overturning streamfunction in the righthand basin. In the model's convention "right" id defined as +eastward/northward and "left" is defined as westward/southward, relative to the model component being discussed. `Nikurashin & Vallis (2012)`_ argued that in the North Atlantic, eastward currents are subducted at the eastern boundary and propagate westward towards the western boundary. @@ -24,8 +25,8 @@ Vanishing net mass flux between the columns yields the boundary conditions: .. math:: \begin{aligned} - \Psi^z_N\left(z=0\right)&=0 \\ - \Psi^z_N\left(z=H\right)&=0 + \Psi\left(z=0\right)&=0 \\ + \Psi\left(z=H\right)&=0 \end{aligned} .. Given a depth :math:`z_o` where :math:`b_N\left(z_o\right)=b_B\left(z_o\right)`, we impose @@ -41,7 +42,7 @@ Vanishing net mass flux between the columns yields the boundary conditions: .. Notice that the above only applies for the equi_colum module, which solves for both the .. overturning streamfunction and buoyancy profles at once - albeit under special conditions. -Finally, the isopycnal overturning streamfunction is obtained by mapping the z-coordinate +The isopycnal overturning streamfunction is obtained by mapping the z-coordinate streamfunction obtained from the thermal wind relation onto the buoyancy in the respective up-stream column: @@ -54,8 +55,8 @@ where :math:`\mathcal{H}` is the Heaviside step function and \begin{aligned} b_{up}\left(z\right) = \begin{cases} - b_N\left(z\right), & \partial_z\Psi\left(z\right) > 0 \\ - b_B\left(z\right), & \partial_z\Psi\left(z\right) < 0 \,. + b_{left}\left(z\right), & \partial_z\Psi\left(z\right) > 0 \\ + b_{right}\left(z\right), & \partial_z\Psi\left(z\right) < 0 \,. \end{cases} \end{aligned} From c1180bf8d8a16c5d74328ad25cd8b107fb104370 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Tue, 25 Feb 2020 07:05:21 -0500 Subject: [PATCH 19/28] More documentation updates --- src/pymoc/modules/psi_thermwind.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/pymoc/modules/psi_thermwind.py b/src/pymoc/modules/psi_thermwind.py index b378281..f1014a4 100644 --- a/src/pymoc/modules/psi_thermwind.py +++ b/src/pymoc/modules/psi_thermwind.py @@ -46,9 +46,9 @@ def __init__( sol_init : ndarray; optional Initial guess at the solution to the thermal wind overturning streamfunction. Units: [...] b1 : float, function, or ndarray; optional - Vertical buoyancy profile from the southern basin. Units: m/s\ :sup:`2` + Vertical buoyancy profile from the righthand (southward/westward) basin. Units: m/s\ :sup:`2` b2 : float, function, or ndarray; optional - Vertical buoyancy profile from the northern basin, representing the + Vertical buoyancy profile from the lefthand (northward/eastward) basin, representing the deepwater formation region. Units: m/s\ :sup:`2` """ @@ -155,8 +155,8 @@ def Psib(self, nb=500): \end{cases} \end{aligned} - where :math:`b_N\left(z\right)` is the density profiles in the northern region, :math:`b_B\left(z\right)` is - the density profile in the southern basin, and :math:`\mathcal{H}` is the Heaviside step function. + where :math:`b_N\left(z\right)` is the density profiles in the righthand region, :math:`b_B\left(z\right)` is + the density profile in the lefthand basin, and :math:`\mathcal{H}` is the Heaviside step function. Parameters ---------- @@ -189,7 +189,7 @@ def Psib(self, nb=500): def Psibz(self, nb=500): r""" Remap the overturning streamfunction onto the native isopycnal-depth space - of the columns in the northern region and southern basin. + of the columns in the right and lefthand basins. Parameters ---------- @@ -199,7 +199,7 @@ def Psibz(self, nb=500): Returns ------- psibz : ndarray - An array where the first element is an array representing the streamfunction at each depth level in the southern basin, and the second represents the same in the northern region. + An array where the first element is an array representing the streamfunction at each depth level in the lefthand basin, and the second represents the same in the righthand basin. """ # map isopycnal overturning back into isopycnal-depth space of each column psib = self.Psib(nb) @@ -216,15 +216,15 @@ def Psibz(self, nb=500): def update(self, b1=None, b2=None): r""" - Update the vertical buoyancy profiles from the southern basin and northern region. + Update the vertical buoyancy profiles from the lefthand basin and righthand basin. Parameters ---------- b1 : float, function, or ndarray; optional - Vertical buoyancy profile from the southern basin. Units: m/s\ :sup:`2` + Vertical buoyancy profile from the lefthand basin. Units: m/s\ :sup:`2` b2 : float, function, or ndarray; optional - Vertical buoyancy profile from the northern basin, representing the + Vertical buoyancy profile from the righthand basin, representing the deepwater formation region. Units: m/s\ :sup:`2` """ # update buoyancy profiles From 54eead4e3e07e5a35c1e0656a668a5a9abe3f60c Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 28 Feb 2020 07:11:07 -0500 Subject: [PATCH 20/28] Fix convection setting --- src/pymoc/modules/column.py | 7 ++----- tests/modules/test_module_wrapper.py | 0 tests/test_model.py | 0 3 files changed, 2 insertions(+), 5 deletions(-) create mode 100644 tests/modules/test_module_wrapper.py create mode 100644 tests/test_model.py diff --git a/src/pymoc/modules/column.py b/src/pymoc/modules/column.py index 6ee57b1..080eb41 100644 --- a/src/pymoc/modules/column.py +++ b/src/pymoc/modules/column.py @@ -231,8 +231,7 @@ def vertadvdiff(self, wA, dt, do_conv=False): dz = self.z[1:] - self.z[:-1] # apply boundary conditions: - # TODO: Make this check the module level do_conv as well - if not self.do_conv: # if we use convection, upper BC is already applied there + if do_conv is None and not self.do_conv or not do_conv: self.b[-1] = self.bs self.b[0] = ( self.bbot if self.bzbot is None else self.b[1] - self.bzbot * dz[0] @@ -340,9 +339,7 @@ def timestep(self, wA=0., dt=1., do_conv=None, vdx_in=None, b_in=None): Buoyancy vales from the adjoining module for the timestepping solution. Units: m/s\ :sup:`2` """ - #TODO: Remove backwards compatibility for method level do_conv from module and examples. Adds unnecesarry complexity. - if self.do_conv: - # if (do_conv is None and self.do_conv) or do_conv: + if do_conv is None and self.do_conv or do_conv: # do convection: (optional) self.convect() diff --git a/tests/modules/test_module_wrapper.py b/tests/modules/test_module_wrapper.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..e69de29 From 6b3907a3fb3d7a3d508fa2702767e4a730258496 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 28 Feb 2020 15:53:08 -0500 Subject: [PATCH 21/28] Stub out model interface tests --- tests/modules/test_module_wrapper.py | 24 ++++++++++++++++++++++++ tests/modules/test_neighbor.py | 14 ++++++++++++++ tests/test_model.py | 28 ++++++++++++++++++++++++++++ 3 files changed, 66 insertions(+) create mode 100644 tests/modules/test_neighbor.py diff --git a/tests/modules/test_module_wrapper.py b/tests/modules/test_module_wrapper.py index e69de29..3b717b4 100644 --- a/tests/modules/test_module_wrapper.py +++ b/tests/modules/test_module_wrapper.py @@ -0,0 +1,24 @@ +import sys +import os +import funcsigs +import numpy as np +from scipy import integrate +import pytest +from pymoc.utils import make_func, make_array +sys.path.append('/pymoc/src/pymoc/modules') +from pymoc.modules import ModuleWrapper + +class TestModuleWrapper(object): + def test_module_wrapper_init(self): + + def test_module_type(self): + + def test_timestep_basin(self): + + def test_update_coupler(self): + + def test_b(self): + + def test_validate_neighbors(self): + + def test_backlink_neighbor_lins(self): \ No newline at end of file diff --git a/tests/modules/test_neighbor.py b/tests/modules/test_neighbor.py new file mode 100644 index 0000000..8df1ccf --- /dev/null +++ b/tests/modules/test_neighbor.py @@ -0,0 +1,14 @@ +import sys +import os +import funcsigs +import numpy as np +from scipy import integrate +import pytest +from pymoc.utils import make_func, make_array +sys.path.append('/pymoc/src/pymoc/modules') +from pymoc.modules import Neighbor + +class TestNeighbor(object): + def test_neighbor_init(self): + + def test_eq(self): diff --git a/tests/test_model.py b/tests/test_model.py index e69de29..8944879 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -0,0 +1,28 @@ +import sys +import os +import funcsigs +import numpy as np +from scipy import integrate +import pytest +from pymoc.utils import make_func, make_array +sys.path.append('/pymoc/src/pymoc/modules') +from pymoc import Model + +class TestModel(object): + def test_model_init(self): + + def test_keys(self): + + def test_get_module(self): + + def test_validate_neighbors_input(self): + + def test_add_module(self): + + def test_new_module(self): + + def test_run(self): + + def test_timestep(self): + + def test_snapshot(self): \ No newline at end of file From f407195180ca000e0bd18bee166f74a69e2d3124 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Wed, 4 Mar 2020 15:42:55 -0500 Subject: [PATCH 22/28] Simplify model neighbors -Remove neighbor class -Move more validation into the module_wrapper -Track explicit left and right neighbor lists -Short circuit backlinking in a friendlier way --- examples/example_twocol_plusSO_model.py | 17 +++-- src/pymoc/model.py | 19 ++---- src/pymoc/modules/__init__.py | 1 - src/pymoc/modules/module_wrapper.py | 85 +++++++++++++------------ src/pymoc/modules/neighbor.py | 10 --- tests/modules/test_module_wrapper.py | 18 +++--- tests/modules/test_neighbor.py | 14 ---- tests/test_model.py | 26 ++++---- 8 files changed, 85 insertions(+), 105 deletions(-) delete mode 100644 src/pymoc/modules/neighbor.py delete mode 100644 tests/modules/test_neighbor.py diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index 2f56309..b53073d 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -9,7 +9,7 @@ and Vallis (2012, JPO). ''' from pymoc import model -from pymoc.modules import Psi_Thermwind, Psi_SO, Column, Neighbor +from pymoc.modules import Psi_Thermwind, Psi_SO, Column from pymoc.plotting import Interpolate_channel import numpy as np from matplotlib import pyplot as plt @@ -87,7 +87,10 @@ def b_north(z): 'bbot': bmin }, 'Atlantic Basin', - neighbors=[Neighbor('psi_so', 'left')] + neighbors=[{ + 'module': model.get_module('psi_so'), + 'direction': 'left' + }] ) # create N.A. overturning model instance model.new_module( @@ -98,7 +101,10 @@ def b_north(z): 'f': 1e-4 }, 'AMOC', - neighbors=[Neighbor('atlantic_basin', 'left')] + neighbors=[{ + 'module': model.get_module('atlantic_basin'), + 'direction': 'left' + }] ) # create adv-diff column model instance for basin @@ -113,7 +119,10 @@ def b_north(z): 'do_conv': True }, 'North Atlantic', - neighbors=[Neighbor('amoc', 'left')] + neighbors=[{ + 'module': model.get_module('amoc'), + 'direction': 'left' + }] ) fig = plt.figure(figsize=(6, 10)) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index c26d873..6febf9d 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -1,5 +1,5 @@ import numpy as np -from pymoc.modules import ModuleWrapper, Neighbor +from pymoc.modules import ModuleWrapper class Model(object): @@ -17,28 +17,17 @@ def get_module(self, key): return self._modules[key] def validate_neighbors_input(self, neighbors): - neighbor_keys = [n.key for n in neighbors] - distinct_neighbor_keys = np.unique(neighbor_keys) - if len(neighbor_keys) > len(distinct_neighbor_keys): + neighbor_modules = [n['module'] for n in neighbors] + distinct_neighbor_moduless = np.unique(neighbor_modules) + if len(neighbor_modules) > len(distinct_neighbor_moduless): raise ValueError( 'Cannot link basins multiple times. Please check your configuration.' ) - for n in neighbors: - neighbor = self.get_module(n.key) - if not neighbor: - raise KeyError( - 'No module present with key ' + n.key + - ', cannot set as neighbor of ' + name - ) - n.module_wrapper = neighbor - def add_module(self, module, name, neighbors=None): neighbors = neighbors or [] self.validate_neighbors_input(neighbors) module_wrapper = ModuleWrapper(module, name, neighbors) - module_wrapper.validate_neighbor_links() - module_wrapper.backlink_neighbors() if hasattr(self, module_wrapper.key): raise NameError( diff --git a/src/pymoc/modules/__init__.py b/src/pymoc/modules/__init__.py index 2048d67..8cdc7eb 100644 --- a/src/pymoc/modules/__init__.py +++ b/src/pymoc/modules/__init__.py @@ -4,5 +4,4 @@ from .psi_SO import Psi_SO from .psi_thermwind import Psi_Thermwind from .SO_ML import SO_ML -from .neighbor import Neighbor from .module_wrapper import ModuleWrapper diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 5ef3e4b..bf7c046 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -1,12 +1,13 @@ import numpy as np -from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column, Neighbor +from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column class ModuleWrapper(object): def __init__(self, module, name, neighbors=None): self.module = module self.name = name - self.neighbors = neighbors or [] + self.left_neighbors = [] + self.right_neighbors = [] self.key = name.replace(' ', '_').lower().strip('_') self.do_psi_bz = hasattr(module, 'Psibz') and callable(module.Psibz) @@ -15,6 +16,9 @@ def __init__(self, module, name, neighbors=None): ) else 'bs' self.psi = [0, 0] + if neighbors: + self.add_neighbors(neighbors) + @property def module_type(self): return self.module.module_type @@ -22,11 +26,10 @@ def module_type(self): def timestep_basin(self, dt=None): module = self.module wA = 0.0 - for neighbor in self.neighbors: - if neighbor.direction == 'right': - wA += neighbor.module_wrapper.psi[0] * 1e6 - else: - wA -= neighbor.module_wrapper.psi[-1] * 1e6 + for neighbor in self.right_neighbors: + wA += neighbor.psi[0] * 1e6 + for neighbor in self.left_neighbors: + wA -= neighbor.psi[-1] * 1e6 module.timestep(wA=wA, dt=dt) @@ -34,11 +37,10 @@ def update_coupler(self): module = self.module b1 = None b2 = None - for neighbor in self.neighbors: - if neighbor.direction == 'left': - b1 = neighbor.module_wrapper and neighbor.module_wrapper.b - else: - b2 = neighbor.module_wrapper and neighbor.module_wrapper.b + for neighbor in self.left_neighbors: + b1 = neighbor.b + for neighbor in self.right_neighbors: + b2 = neighbor.b if self.do_psi_bz: module.update(b1=b1, b2=b2) @@ -51,37 +53,42 @@ def update_coupler(self): def b(self): return getattr(self.module, self.b_type) - def validate_neighbors(self, backlink=False): - if not backlink: - for neighbor in self.neighbors: - linked_neighbor_keys = [ - n.key for n in neighbor.module_wrapper.neighbors - ] - if self.key in linked_neighbor_keys: - raise KeyError( - 'Cannot add module ' + self.name + ' as a neighbor of ' + - neighbor.module_wrapper.name + - ' because they are already coupled.' - ) + @property + def neighbors(self): + return self.left_neighbors + self.right_neighbors + + def add_neighbor(self, new_neighbor, direction, backlink=False): + if not direction in ['left', 'right']: + raise ValueError( + "Direction of a neighbor must be either 'left' or 'right'." + ) + + if new_neighbor in self.neighbors: + raise KeyError( + 'Cannot add module ' + new_neighbor.name + ' as a neighbor of ' + + self.name + ' because they are already coupled.' + ) if self.module_type == 'coupler': # We don't need to explicitly check that a coupler has two or fewer neighbors, as the enforcement of only one left, only one right, and only being able to point left or right implicitly enforces that condition. - left_neighbor_count = sum(n.direction == 'left' for n in self.neighbors) - right_neighbor_count = sum( - n.direction == 'right' for n in self.neighbors - ) - if left_neighbor_count > 1 or right_neighbor_count > 1: + if len(self.left_neighbors) > 1 or len(self.right_neighbors) > 1: raise ValueError( 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' ) - def backlink_neighbor_links(self): - for neighbor in self.neighbors: - neighbor.module_wrapper.neighbors.append( - Neighbor( - self.key, - 'left' if neighbor.direction == 'right' else 'right', - module_wrapper=self - ) - ) - neighbor.module_wrapper.validate_neighbors(backlink=True) + if direction == 'left': + self.left_neighbors.append(new_neighbor) + else: + self.right_neighbors.append(new_neighbor) + + if not backlink: + self.backlink_neighbor(new_neighbor) + + def add_neighbors(self, neighbors): + if len(neighbors) > 0: + for neighbor in neighbors: + self.add_neighbor(neighbor['module'], neighbor['direction']) + + def backlink_neighbor(self, neighbor): + direction = 'left' if neighbor in self.right_neighbors else 'right' + neighbor.add_neighbor(self, direction, backlink=True) diff --git a/src/pymoc/modules/neighbor.py b/src/pymoc/modules/neighbor.py deleted file mode 100644 index 5044618..0000000 --- a/src/pymoc/modules/neighbor.py +++ /dev/null @@ -1,10 +0,0 @@ -class Neighbor(object): - def __init__(self, key, direction, module_wrapper=None): - self.key = key - if direction not in ['left', 'right']: - raise ValueError('Direction must be either left or right.') - self.direction = direction - self.module_wrapper = module_wrapper - - def __eq__(self, other): - return self.key == other.key and self.direction == other.direction diff --git a/tests/modules/test_module_wrapper.py b/tests/modules/test_module_wrapper.py index 3b717b4..1e494b0 100644 --- a/tests/modules/test_module_wrapper.py +++ b/tests/modules/test_module_wrapper.py @@ -8,17 +8,17 @@ sys.path.append('/pymoc/src/pymoc/modules') from pymoc.modules import ModuleWrapper -class TestModuleWrapper(object): - def test_module_wrapper_init(self): +# class TestModuleWrapper(object): +# def test_module_wrapper_init(self): - def test_module_type(self): + # def test_module_type(self): - def test_timestep_basin(self): - - def test_update_coupler(self): + # def test_timestep_basin(self): - def test_b(self): + # def test_update_coupler(self): - def test_validate_neighbors(self): + # def test_b(self): - def test_backlink_neighbor_lins(self): \ No newline at end of file + # def test_validate_neighbors(self): + + # def test_backlink_neighbor_lins(self): diff --git a/tests/modules/test_neighbor.py b/tests/modules/test_neighbor.py deleted file mode 100644 index 8df1ccf..0000000 --- a/tests/modules/test_neighbor.py +++ /dev/null @@ -1,14 +0,0 @@ -import sys -import os -import funcsigs -import numpy as np -from scipy import integrate -import pytest -from pymoc.utils import make_func, make_array -sys.path.append('/pymoc/src/pymoc/modules') -from pymoc.modules import Neighbor - -class TestNeighbor(object): - def test_neighbor_init(self): - - def test_eq(self): diff --git a/tests/test_model.py b/tests/test_model.py index 8944879..bc7d263 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -5,24 +5,24 @@ from scipy import integrate import pytest from pymoc.utils import make_func, make_array -sys.path.append('/pymoc/src/pymoc/modules') -from pymoc import Model +# sys.path.append('/pymoc/src/pymoc') +# from pymoc import Model -class TestModel(object): - def test_model_init(self): + # class TestModel(object): + # def test_model_init(self): - def test_keys(self): + # def test_keys(self): - def test_get_module(self): - - def test_validate_neighbors_input(self): + # def test_get_module(self): - def test_add_module(self): + # def test_validate_neighbors_input(self): - def test_new_module(self): + # def test_add_module(self): - def test_run(self): + # def test_new_module(self): - def test_timestep(self): + # def test_run(self): - def test_snapshot(self): \ No newline at end of file + # def test_timestep(self): + + # def test_snapshot(self): From 246bf4ae5e062060b001aab3b7bc550c66470987 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Fri, 6 Mar 2020 07:27:05 -0500 Subject: [PATCH 23/28] minor refactoring of module_wrapper --- src/pymoc/modules/module_wrapper.py | 72 ++++++++++++++++++----------- 1 file changed, 46 insertions(+), 26 deletions(-) diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index bf7c046..4165cb3 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -35,17 +35,12 @@ def timestep_basin(self, dt=None): def update_coupler(self): module = self.module - b1 = None - b2 = None - for neighbor in self.left_neighbors: - b1 = neighbor.b - for neighbor in self.right_neighbors: - b2 = neighbor.b if self.do_psi_bz: - module.update(b1=b1, b2=b2) + module.update(b1=self.b1, b2=self.b2) else: - module.update(b=b2) + module.update(b=self.b2) + module.solve() self.psi = module.Psibz() if self.do_psi_bz else [module.Psi] @@ -57,38 +52,63 @@ def b(self): def neighbors(self): return self.left_neighbors + self.right_neighbors - def add_neighbor(self, new_neighbor, direction, backlink=False): + @property + def b1(self): + if self.module_type != 'coupler': + raise TypeError('Cannot access b1 for non-coupler modules.') + if len(left_neighbors) > 0: + return self.left_neighbors[0].b + return None + + @property + def b2(self): + if self.module_type != 'coupler': + raise TypeError('Cannot access b2 for non-coupler modules.') + if len(right_neighbors) > 0: + return self.right_neighbors[0].b + return None + + def add_neighbor(self, new_neighbor, direction, backlinking=False): + self.validate_neighbor_direction(direction) + self.validate_neighbor_uniqueness(new_neighbor) + self.validate_coupler_neighbor_direction(direction) + + if direction == 'left': + self.left_neighbors.append(new_neighbor) + else: + self.right_neighbors.append(new_neighbor) + + if not backlinklinking: + self.backlink_neighbor(new_neighbor) + + def add_neighbors(self, neighbors): + if len(neighbors) > 0: + for neighbor in neighbors: + self.add_neighbor(neighbor['module'], neighbor['direction']) + + def validate_neighbor_direction(self, direction): if not direction in ['left', 'right']: raise ValueError( "Direction of a neighbor must be either 'left' or 'right'." ) - if new_neighbor in self.neighbors: + def validate_neighbor_uniqueness(self, neighbor): + if neighbor in self.neighbors: raise KeyError( - 'Cannot add module ' + new_neighbor.name + ' as a neighbor of ' + + 'Cannot add module ' + neighbor.name + ' as a neighbor of ' + self.name + ' because they are already coupled.' ) + def validate_coupler_neighbor_direction(self, direction): if self.module_type == 'coupler': # We don't need to explicitly check that a coupler has two or fewer neighbors, as the enforcement of only one left, only one right, and only being able to point left or right implicitly enforces that condition. - if len(self.left_neighbors) > 1 or len(self.right_neighbors) > 1: + if len(self.left_neighbors) > 0 and direction == 'left' or len( + self.right_neighbors + ) > 0 and direction == 'right': raise ValueError( 'Cannot have a coupler linked in the same direction more than once. Please check your configuration.' ) - if direction == 'left': - self.left_neighbors.append(new_neighbor) - else: - self.right_neighbors.append(new_neighbor) - - if not backlink: - self.backlink_neighbor(new_neighbor) - - def add_neighbors(self, neighbors): - if len(neighbors) > 0: - for neighbor in neighbors: - self.add_neighbor(neighbor['module'], neighbor['direction']) - def backlink_neighbor(self, neighbor): direction = 'left' if neighbor in self.right_neighbors else 'right' - neighbor.add_neighbor(self, direction, backlink=True) + neighbor.add_neighbor(self, direction, backlinking=True) From 9cd5b16f3b4cd2c31243045e02e6365f2c87831c Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Mon, 9 Mar 2020 06:44:11 -0400 Subject: [PATCH 24/28] more neighbor cleanup, make code more modular and DRY --- examples/twobasin_NadeauJansen_model.py | 28 +++++++++++++++++-------- src/pymoc/modules/module_wrapper.py | 6 +++--- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/examples/twobasin_NadeauJansen_model.py b/examples/twobasin_NadeauJansen_model.py index d835b37..e363ac5 100644 --- a/examples/twobasin_NadeauJansen_model.py +++ b/examples/twobasin_NadeauJansen_model.py @@ -13,7 +13,7 @@ ''' from pymoc import model -from pymoc.modules import Psi_Thermwind, Psi_SO, Column, Neighbor +from pymoc.modules import Psi_Thermwind, Psi_SO, Column from pymoc.plotting import Interpolate_channel, Interpolate_twocol import numpy as np from matplotlib import pyplot as plt @@ -133,11 +133,14 @@ def b_Pac(z): 'N2min': N2min }, 'Atl', - neighbors=[ - Neighbor('amoc', 'right'), - Neighbor('so_atl', 'left'), - Neighbor('zoc', 'right') - ] + neighbors=[{ + 'module': m.get_module('amoc'), + 'direction': 'right', + 'module': m.get_module('so_atl'), + 'direction': 'left', + 'module': m.get_module('zoc'), + 'direction': 'right' + }] ) # create adv-diff column model instance for northern sinking region @@ -153,7 +156,10 @@ def b_Pac(z): 'do_conv': True }, 'North', - neighbors=[Neighbor('amoc', 'left')] + neighbors=[{ + 'module': m.get_module('amoc'), + 'direction': 'left' + }] ) # create adv-diff column model instance for Pac @@ -168,8 +174,12 @@ def b_Pac(z): 'N2min': N2min }, 'Pac', - neighbors=[Neighbor('zoc', 'left'), - Neighbor('so_pac', 'left')] + neighbors=[{ + 'module': m.get_module('zoc'), + 'direction': 'left', + 'module': m.get_module('so_pac'), + 'direction': 'left' + }] ) # Create figure: diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 4165cb3..15016e1 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -56,7 +56,7 @@ def neighbors(self): def b1(self): if self.module_type != 'coupler': raise TypeError('Cannot access b1 for non-coupler modules.') - if len(left_neighbors) > 0: + if len(self.left_neighbors) > 0: return self.left_neighbors[0].b return None @@ -64,7 +64,7 @@ def b1(self): def b2(self): if self.module_type != 'coupler': raise TypeError('Cannot access b2 for non-coupler modules.') - if len(right_neighbors) > 0: + if len(self.right_neighbors) > 0: return self.right_neighbors[0].b return None @@ -78,7 +78,7 @@ def add_neighbor(self, new_neighbor, direction, backlinking=False): else: self.right_neighbors.append(new_neighbor) - if not backlinklinking: + if not backlinking: self.backlink_neighbor(new_neighbor) def add_neighbors(self, neighbors): From 7bdd2b879784edd6aa40bd052f9a1c8e0eab5862 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Wed, 11 Mar 2020 06:55:41 -0400 Subject: [PATCH 25/28] Simplify neighbor data model --- examples/example_twocol_plusSO_model.py | 24 ++++----- src/pymoc/model.py | 27 ++++++---- src/pymoc/modules/module_wrapper.py | 72 +++++++++++-------------- 3 files changed, 59 insertions(+), 64 deletions(-) diff --git a/examples/example_twocol_plusSO_model.py b/examples/example_twocol_plusSO_model.py index b53073d..c170bb7 100644 --- a/examples/example_twocol_plusSO_model.py +++ b/examples/example_twocol_plusSO_model.py @@ -78,7 +78,8 @@ def b_north(z): ) # create adv-diff column model instance for basin model.new_module( - Column, { + Column, + { 'z': z, 'kappa': kappa, 'Area': A_basin, @@ -87,29 +88,25 @@ def b_north(z): 'bbot': bmin }, 'Atlantic Basin', - neighbors=[{ - 'module': model.get_module('psi_so'), - 'direction': 'left' - }] + left_neighbors=[model.get_module('psi_so')], ) # create N.A. overturning model instance model.new_module( - Psi_Thermwind, { + Psi_Thermwind, + { 'z': z, 'b1': b_basin, 'b2': b_north, 'f': 1e-4 }, 'AMOC', - neighbors=[{ - 'module': model.get_module('atlantic_basin'), - 'direction': 'left' - }] + left_neighbors=[model.get_module('atlantic_basin')], ) # create adv-diff column model instance for basin model.new_module( - Column, { + Column, + { 'z': z, 'kappa': kappa, 'Area': A_north, @@ -119,10 +116,7 @@ def b_north(z): 'do_conv': True }, 'North Atlantic', - neighbors=[{ - 'module': model.get_module('amoc'), - 'direction': 'left' - }] + left_neighbors=[model.get_module('amoc')], ) fig = plt.figure(figsize=(6, 10)) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index 6febf9d..edf4052 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -8,6 +8,9 @@ def __init__(self): self.couplers = [] self._modules = {} + def __getitem__(self, key): + return self._modules[key] + def keys(self): return self._modules.keys() @@ -17,17 +20,21 @@ def get_module(self, key): return self._modules[key] def validate_neighbors_input(self, neighbors): - neighbor_modules = [n['module'] for n in neighbors] - distinct_neighbor_moduless = np.unique(neighbor_modules) - if len(neighbor_modules) > len(distinct_neighbor_moduless): + neighbor_modules = [['module'] for n in neighbors] + distinct_neighbor_modules = np.unique(neighbor_modules) + if len(neighbor_modules) > len(distinct_neighbor_modules): raise ValueError( 'Cannot link basins multiple times. Please check your configuration.' ) - def add_module(self, module, name, neighbors=None): - neighbors = neighbors or [] - self.validate_neighbors_input(neighbors) - module_wrapper = ModuleWrapper(module, name, neighbors) + def add_module(self, module, name, left_neighbors=[], right_neighbors=[]): + self.validate_neighbors_input(left_neighbors + right_neighbors) + module_wrapper = ModuleWrapper( + module, + name, + left_neighbors=left_neighbors, + right_neighbors=right_neighbors + ) if hasattr(self, module_wrapper.key): raise NameError( @@ -48,12 +55,14 @@ def new_module( module_class, module_args, module_name, - neighbors=[], + left_neighbors=[], + right_neighbors=[] ): self.add_module( module_class(**module_args), module_name, - neighbors, + left_neighbors=left_neighbors, + right_neighbors=right_neighbors ) def run( diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 15016e1..264a556 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -3,7 +3,7 @@ class ModuleWrapper(object): - def __init__(self, module, name, neighbors=None): + def __init__(self, module, name, left_neighbors=None, right_neighbors=None): self.module = module self.name = name self.left_neighbors = [] @@ -16,8 +16,10 @@ def __init__(self, module, name, neighbors=None): ) else 'bs' self.psi = [0, 0] - if neighbors: - self.add_neighbors(neighbors) + if left_neighbors or right_neighbors: + self.add_neighbors( + left_neighbors=left_neighbors, right_neighbors=right_neighbors + ) @property def module_type(self): @@ -34,12 +36,17 @@ def timestep_basin(self, dt=None): module.timestep(wA=wA, dt=dt) def update_coupler(self): + if self.module_type != 'coupler': + raise TypeError('Cannot use update_coupler on non-coupler modules.') + module = self.module + b1 = self.left_neighbors[0].b if len(self.left_neighbors) > 0 else None + b2 = self.right_neighbors[0].b if len(self.right_neighbors) > 0 else None if self.do_psi_bz: - module.update(b1=self.b1, b2=self.b2) + module.update(b1=b1, b2=b2) else: - module.update(b=self.b2) + module.update(b=b2) module.solve() self.psi = module.Psibz() if self.do_psi_bz else [module.Psi] @@ -52,45 +59,28 @@ def b(self): def neighbors(self): return self.left_neighbors + self.right_neighbors - @property - def b1(self): - if self.module_type != 'coupler': - raise TypeError('Cannot access b1 for non-coupler modules.') - if len(self.left_neighbors) > 0: - return self.left_neighbors[0].b - return None - - @property - def b2(self): - if self.module_type != 'coupler': - raise TypeError('Cannot access b2 for non-coupler modules.') - if len(self.right_neighbors) > 0: - return self.right_neighbors[0].b - return None - - def add_neighbor(self, new_neighbor, direction, backlinking=False): - self.validate_neighbor_direction(direction) + def add_left_neighbor(self, new_neighbor, backlinking=False): self.validate_neighbor_uniqueness(new_neighbor) - self.validate_coupler_neighbor_direction(direction) - - if direction == 'left': - self.left_neighbors.append(new_neighbor) - else: - self.right_neighbors.append(new_neighbor) + self.validate_coupler_neighbor_direction('left') + self.left_neighbors.append(new_neighbor) + if not backlinking: + self.backlink_neighbor(new_neighbor) + def add_right_neighbor(self, new_neighbor, backlinking=False): + self.validate_neighbor_uniqueness(new_neighbor) + self.validate_coupler_neighbor_direction('right') + self.right_neighbors.append(new_neighbor) if not backlinking: self.backlink_neighbor(new_neighbor) - def add_neighbors(self, neighbors): - if len(neighbors) > 0: - for neighbor in neighbors: - self.add_neighbor(neighbor['module'], neighbor['direction']) + def add_neighbors(self, left_neighbors=None, right_neighbors=None): + if len(left_neighbors) > 0: + for neighbor in left_neighbors: + self.add_left_neighbor(neighbor) - def validate_neighbor_direction(self, direction): - if not direction in ['left', 'right']: - raise ValueError( - "Direction of a neighbor must be either 'left' or 'right'." - ) + if len(right_neighbors) > 0: + for neighbor in right_neighbors: + self.add_right_neighbor(neighbor) def validate_neighbor_uniqueness(self, neighbor): if neighbor in self.neighbors: @@ -110,5 +100,7 @@ def validate_coupler_neighbor_direction(self, direction): ) def backlink_neighbor(self, neighbor): - direction = 'left' if neighbor in self.right_neighbors else 'right' - neighbor.add_neighbor(self, direction, backlinking=True) + if neighbor in self.right_neighbors: + neighbor.add_left_neighbor(self, backlinking=True) + else: + neighbor.add_right_neighbor(self, backlinking=True) From 6704a3fc12dd839d15e6814b7e5ad67e4b350549 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sat, 21 Mar 2020 13:45:12 -0400 Subject: [PATCH 26/28] Begin module_wrapper documentation --- src/pymoc/modules/module_wrapper.py | 41 +++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 264a556..4867163 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -1,9 +1,46 @@ import numpy as np -from pymoc.modules import Psi_SO, Psi_Thermwind, SO_ML, Column, Equi_Column +from pymoc.modules import Column, Equi_Column class ModuleWrapper(object): - def __init__(self, module, name, left_neighbors=None, right_neighbors=None): + r""" + Module Wrapper + + Instances of this class wrap individual phsyical modules + (e.g. advective-diffusive columns, thermal wind streamfunctions) + for use in the model driver. The module wrapper is responsible for + timestepping & updating modules, & communication between neighboring + modules. + + By convention geographic north & east are defined as + being to the "right" or "rightward", while geographic south & west + are defined as being the the "left" or "leftward". Modules are categorized + as "basins" with time evolving density profiles (e.g. columns, mixed layers), + or "couplers" with streamfunctions that maintain the dynamical balance between + neighboring basins (e.g. thermal wind streamfunction, SO transport). + """ + def __init__( + self, + module, + name, + left_neighbors=None, + right_neighbors=None, + ): + r""" + Parameters + ---------- + + module : module class instance + Physics module being wrapped + name : string + Name of the module being wrapped (e.g. Atlantic Basin, AMOC) + left_neighbors : list; optional + List of modules to the "left" of the module being wrapped. + Couplers may have at most one left neighbor. + right_neighbors : list; optional + List of modules to the "right" of the module being wrapped. + Couplers may have at most one right neighbor. + """ self.module = module self.name = name self.left_neighbors = [] From b37779ac1bf55a71901ded58eafc4980af94c20e Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Tue, 21 Apr 2020 17:57:41 -0400 Subject: [PATCH 27/28] Document the module_wrapper --- src/pymoc/modules/module_wrapper.py | 127 +++++++++++++++++++++++++++- 1 file changed, 125 insertions(+), 2 deletions(-) diff --git a/src/pymoc/modules/module_wrapper.py b/src/pymoc/modules/module_wrapper.py index 4867163..676c105 100644 --- a/src/pymoc/modules/module_wrapper.py +++ b/src/pymoc/modules/module_wrapper.py @@ -60,9 +60,28 @@ def __init__( @property def module_type(self): + r""" + Return whether the wrapped module is a basin or coupler. + + Returns + ------- + + module_type : string + The type of the wrapped module. + """ return self.module.module_type def timestep_basin(self, dt=None): + r""" + Invoke the timestep method of the wrapped basin module. + + Parameters + ---------- + + dt : int + Numerical timestep over which solution are iterated. Units: s + + """ module = self.module wA = 0.0 for neighbor in self.right_neighbors: @@ -73,6 +92,10 @@ def timestep_basin(self, dt=None): module.timestep(wA=wA, dt=dt) def update_coupler(self): + r""" + Invoke the update method of the wrapped coupler. + + """ if self.module_type != 'coupler': raise TypeError('Cannot use update_coupler on non-coupler modules.') @@ -90,13 +113,52 @@ def update_coupler(self): @property def b(self): + r""" + Return the buoyancy profile of the wrapped module. This is the + vertical buoyancy profile for columns, and the surface buoyancy + for channels and couplers + + Returns + ------- + + module_type : func or ndarray + The buoyancy profile of the wrapped module, with type consistent + with the module. + """ return getattr(self.module, self.b_type) @property def neighbors(self): + r""" + Return a list of all of the module's neighbors. + + Returns + ------- + + neighbors : ndarray + Array pointing to the wrappers for each + neighboring module. + + """ return self.left_neighbors + self.right_neighbors def add_left_neighbor(self, new_neighbor, backlinking=False): + r""" + Add a neighbor to the "left" of the current module. This method validates + that the neighbor is unique, can occupy the left neighbor position, and + optionally backlinks to the current module (i.e. sets the current module + as the new neighbor's right neighbor). + + + Parameters + ---------- + + new_neighbor: module_wrapper + A wrapper pointing to the module being added as a lefthand neighbor. + backlinking: boolean; optional + Whether to backlink the new neighbor to the current module. Defaults false. + + """ self.validate_neighbor_uniqueness(new_neighbor) self.validate_coupler_neighbor_direction('left') self.left_neighbors.append(new_neighbor) @@ -104,6 +166,22 @@ def add_left_neighbor(self, new_neighbor, backlinking=False): self.backlink_neighbor(new_neighbor) def add_right_neighbor(self, new_neighbor, backlinking=False): + r""" + Add a neighbor to the "right" of the current module. This method validates + that the neighbor is unique, can occupy the right neighbor position, and + optionally backlinks to the current module (i.e. sets the current module + as the new neighbor's left neighbor). + + + Parameters + ---------- + + new_neighbor: module_wrapper + A wrapper pointing to the module being added as a righthand neighbor. + backlinking: boolean; optional + Whether to backlink the new neighbor to the current module. Defaults false. + + """ self.validate_neighbor_uniqueness(new_neighbor) self.validate_coupler_neighbor_direction('right') self.right_neighbors.append(new_neighbor) @@ -111,15 +189,38 @@ def add_right_neighbor(self, new_neighbor, backlinking=False): self.backlink_neighbor(new_neighbor) def add_neighbors(self, left_neighbors=None, right_neighbors=None): - if len(left_neighbors) > 0: + r""" + Add multiple neighbors to the current module. + + Parameters + ---------- + + left_neighbors: ndarray; optional + A list of modules to be added to the left of the current module + right_neighbors: ndarray; optional + A list of modules to be added to the right of the current module + + """ + if left_neighbors is not None and len(left_neighbors) > 0: for neighbor in left_neighbors: self.add_left_neighbor(neighbor) - if len(right_neighbors) > 0: + if right_neighbors is not None and len(right_neighbors) > 0: for neighbor in right_neighbors: self.add_right_neighbor(neighbor) def validate_neighbor_uniqueness(self, neighbor): + r""" + Validate that the current module does not already have the specified module as a neighbor. + If already a neighbor, raises a KeyError exception. + + Parameters + ---------- + + neighbor: module_wrapper + The module being checked against for uniqueness + + """ if neighbor in self.neighbors: raise KeyError( 'Cannot add module ' + neighbor.name + ' as a neighbor of ' + @@ -127,6 +228,17 @@ def validate_neighbor_uniqueness(self, neighbor): ) def validate_coupler_neighbor_direction(self, direction): + r""" + Validate that the current module does not already have a neighbor in the specified direction + if a coupler. If the current module is a coupler, and direction is occupied, raises a ValueError exception. + + Parameters + ---------- + + direction: string + The direction being checked against for availability + + """ if self.module_type == 'coupler': # We don't need to explicitly check that a coupler has two or fewer neighbors, as the enforcement of only one left, only one right, and only being able to point left or right implicitly enforces that condition. if len(self.left_neighbors) > 0 and direction == 'left' or len( @@ -137,6 +249,17 @@ def validate_coupler_neighbor_direction(self, direction): ) def backlink_neighbor(self, neighbor): + r""" + Point the specified neighbor back at the current module. If a right neighbor, the current module becomes + its left neighbor. Otherwise, the current module becomes its right neighbor. + + Parameters + ---------- + + neighbor: module_wrapper + Wrapper pointing to the module being backlinked from + + """ if neighbor in self.right_neighbors: neighbor.add_left_neighbor(self, backlinking=True) else: From 58686f6cb2c22330882c31458b0e704ad28ef173 Mon Sep 17 00:00:00 2001 From: Mike Bueti Date: Sat, 9 May 2020 12:09:18 -0400 Subject: [PATCH 28/28] Document the Model class --- src/pymoc/model.py | 148 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 146 insertions(+), 2 deletions(-) diff --git a/src/pymoc/model.py b/src/pymoc/model.py index edf4052..44289f6 100644 --- a/src/pymoc/model.py +++ b/src/pymoc/model.py @@ -3,6 +3,15 @@ class Model(object): + r""" + Model + + Instances of this class maintain a collection of physics modules, + their relationships to one another, and their evolution in time. + Creation of a Model instance allows users to integrate arbitrary + basin configurations without needing to manually track, timestep, + and reconcile the individual modules. + """ def __init__(self): self.basins = [] self.couplers = [] @@ -12,15 +21,50 @@ def __getitem__(self, key): return self._modules[key] def keys(self): + r""" + Get a list of all module keys for the model configuration. + + Returns + ------- + + keys : list + A list containing key strings for each module. + """ return self._modules.keys() def get_module(self, key): + r""" + Retrieve a model module by its key. + + Parameters + ---------- + + key : string + Key of the module being retrieved + + Returns + ------- + + module : ModuleWrapperr + ModuleWrapper instance associated with the specified key + """ if not key: return None return self._modules[key] def validate_neighbors_input(self, neighbors): - neighbor_modules = [['module'] for n in neighbors] + r""" + Ensure that a set of neighbors is valid, specifically that no neighbors are being + linked multiple times at once. Raises a ValueError if validation fails. + + Parameters + ---------- + + neighbors : list + A list of ModuleWrappers pointing a the modules to be validated + """ + + neighbor_modules = [n['module'] for n in neighbors] distinct_neighbor_modules = np.unique(neighbor_modules) if len(neighbor_modules) > len(distinct_neighbor_modules): raise ValueError( @@ -28,6 +72,26 @@ def validate_neighbors_input(self, neighbors): ) def add_module(self, module, name, left_neighbors=[], right_neighbors=[]): + r""" + Add a physical module to the model configuration, with its location determined + by the specified neighbors. + + Parameters + ---------- + + module : Module instance + The physics module (e.g. Column, PsiThermwind, etc.) to be added to the + model configuration + name : string + The name with which to label the module in the model configuration + left_neighbors : list + List of existing model modules that will be to the "left" of the + newly added module + right_neighbors : list + List of existing model modules that will be to the "right" of the + newly added module + """ + self.validate_neighbors_input(left_neighbors + right_neighbors) module_wrapper = ModuleWrapper( module, @@ -58,6 +122,28 @@ def new_module( left_neighbors=[], right_neighbors=[] ): + r""" + Create a new physical module and add itto the model configuration, with its + location determined by the specified neighbors. + + Parameters + ---------- + + module_class : Class + The class constructor of the physics module to be created and added + to the model configuration + moduel_args : dict + A dictionary of argument keys and values to be used to initialize + the newly created module + module_name : string + The name with which to label the module in the model configuration + left_neighbors : list + List of existing model modules that will be to the "left" of the + newly added module + right_neighbors : list + List of existing model modules that will be to the "right" of the + newly added module + """ self.add_module( module_class(**module_args), module_name, @@ -69,10 +155,31 @@ def run( self, steps, basin_dt, - coupler_dt, + coupler_dt=0, snapshot_start=None, snapshot_interval=None ): + r""" + Integrate the model forward in time. This method will timestep all basin modules at the + basin timestep, and update all coupler modules at the coupling timestep. This method will + optionally yield, allowing diagnosis of the model's temporal evolution. + + Parameters + ---------- + + steps: int + The number of model integration timesteps + basin_dt : int + The timestep length (in seconds) for basin modules + coupler_dt : int, optional + The duration (in seconds) between coupler module updates. If unspecified, + coupling takes place at every timestep + snapsot_start : int, optional + The model step at which to begin yielding for model snapshotting + snapshot_interval : int, optional + The number of steps between yields for model snapshotting + """ + for coupler in self.couplers: coupler.update_coupler() @@ -83,6 +190,22 @@ def run( return def timestep(self, step, basin_dt, coupler_dt=0): + r""" + Execute a single model integration step, which timesteps all basins and updates + all couplers at the coupling timestep + + Parameters + ---------- + + step : int + The current model timestep + basin_dt : int + The timestep length (in seconds) for basin modules + coupler_dt : int, optional + The duration (in seconds) between coupler module updates. If unspecified, + coupling takes place at every timestep + """ + for basin in self.basins: basin.timestep_basin(dt=basin_dt) if step % coupler_dt == 0: @@ -90,6 +213,27 @@ def timestep(self, step, basin_dt, coupler_dt=0): coupler.update_coupler() def snapshot(self, step, snapshot_start, snapshot_interval): + r""" + Determine whether the current model timestep should yield to snapshotting. + + Parameters + ---------- + + step : int + The current model timestep + snapsot_start : int + The model step at which to begin yielding for model snapshotting + snapshot_interval : int + The number of steps between yields for model snapshotting + + Returns + ------- + + snapshot : bool + Whether all criteria are met for the model to yield to snapshotting on + the current timestep + """ + return snapshot_start is not None and snapshot_interval is not None and step >= snapshot_start and ( step-snapshot_start ) % snapshot_interval == 0