Add callback functionality

pipedream_solver/callbacks.py

@@ -0,0 +1,22 @@
+class BaseCallback():
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __on_step_start__(self, *args, **kwargs):
+        return None
+
+    def __on_step_end__(self, *args, **kwargs):
+        return None
+
+    def __on_save_state__(self, *args, **kwargs):
+        return None
+
+    def __on_load_state__(self, *args, **kwargs):
+        return None
+
+    def __on_simulation_start__(self, *args, **kwargs):
+        return None
+
+    def __on_simulation_end__(self, *args, **kwargs):
+        return None
+

pipedream_solver/convergence.py

@@ -0,0 +1,70 @@
+import numpy as np
+from pipedream_solver.callbacks import BaseCallback
+
+class BaseConvergenceManager(BaseCallback):
+    def __init__(self, model, num_iter):
+        self.model = model
+        self.num_iter = num_iter
+        self.iter_elapsed = 0
+
+    def __on_step_start__(self, *args, **kwargs):
+        if (self.iter_elapsed == 0):
+            self.model.save_state()
+
+    def __on_step_end__(self, *args, **kwargs):
+        dt = kwargs['dt']
+        self.iter_elapsed += 1
+        while (self.iter_elapsed < self.num_iter):
+            condition_met = self.convergence_condition()
+            if not condition_met:
+                self.H_j_prev = self.model.H_j.copy()
+                self.model.iter_count -= 1
+                self.model.t -= dt
+                self.model._setup_step(*args, **kwargs)
+                self.model._solve_step(*args, **kwargs)
+                self.H_j_next = self.model.H_j.copy()
+                self.iter_elapsed += 1
+            else:
+                break
+        self.iter_elapsed = 0
+
+    def convergence_condition(self):
+        return False
+
+class DefaultConvergenceManager(BaseConvergenceManager):
+    def __init__(self, model, head_tol, num_iter):
+        self.model = model
+        self.head_tol = head_tol
+        self.num_iter = num_iter
+        self.iter_elapsed = 0
+
+    def __on_step_start__(self, *args, **kwargs):
+        self.H_j_prev = self.model.H_j.copy()
+        if (self.iter_elapsed == 0):
+            self.model.save_state()
+
+    def __on_step_end__(self, *args, **kwargs):
+        dt = kwargs['dt']
+        self.H_j_next = self.model.H_j.copy()
+        self.iter_elapsed += 1
+        while (self.iter_elapsed < self.num_iter):
+            condition_met = self.convergence_condition()
+            if not condition_met:
+                self.H_j_prev = self.model.H_j.copy()
+                self.model.iter_count -= 1
+                self.model.t -= dt
+                self.model._setup_step(*args, **kwargs)
+                self.model._solve_step(*args, **kwargs)
+                self.H_j_next = self.model.H_j.copy()
+                self.iter_elapsed += 1
+            else:
+                break
+        self.iter_elapsed = 0
+
+    def convergence_condition(self):
+        head_tol = self.head_tol
+        H_j_prev = self.H_j_prev
+        H_j_next = self.H_j_next
+        residual = np.abs(H_j_next - H_j_prev)
+        condition_met = (residual < head_tol).all()
+        return condition_met

pipedream_solver/diagnostics.py

@@ -0,0 +1,157 @@
+import numpy as np
+from pipedream_solver._nsuperlink import numba_compute_functional_storage_volumes, numba_compute_tabular_storage_volumes
+from pipedream_solver.callbacks import BaseCallback
+
+class ErrorTracker(BaseCallback):
+    def __init__(self, model):
+        self.model = model
+        self.error = None
+
+    def __on_step_end__(self):
+        model = self.model
+        dt = model._dt
+        error = compute_error(model, dt)
+        self.error = error
+
+
+def continuity_error_j(model, dt):
+    error = np.zeros(model.M)
+    H_j_next = model.H_j
+    H_j_prev = model.states['H_j']
+    bc = model.bc
+    Q_in = model._Q_in
+    error += model.A_sj / dt
+    np.add.at(error, model._J_uk, model._B_uk * model._dx_uk * model._theta_uk / 2 / dt)
+    np.add.at(error, model._J_dk, model._B_dk * model._dx_dk * model._theta_dk / 2 / dt)
+    error *= (H_j_next - H_j_prev)
+    np.add.at(error, model._J_uk, model._Q_uk)
+    np.subtract.at(error, model._J_dk, model._Q_dk)
+    np.add.at(error, model._J_uo, model._Qo)
+    np.subtract.at(error, model._J_do, model._Qo)
+    np.add.at(error, model._J_uw, model._Qw)
+    np.subtract.at(error, model._J_dw, model._Qw)
+    np.add.at(error, model._J_up, model._Qp)
+    np.subtract.at(error, model._J_dp, model._Qp)
+    error -= Q_in
+    error[bc] = 0.
+    return error
+
+def momentum_error_uk(model, dt):
+    error = np.zeros(model.NK)
+    raise NotImplementedError
+
+def momentum_error_dk(model, dt):
+    error = np.zeros(model.NK)
+    raise NotImplementedError
+
+def momentum_error_o(model, dt):
+    error = np.zeros(model.n_o)
+    raise NotImplementedError
+
+def momentum_error_w(model, dt):
+    error = np.zeros(model.n_w)
+    raise NotImplementedError
+
+def momentum_error_p(model, dt):
+    error = np.zeros(model.n_p)
+    raise NotImplementedError
+
+def continuity_error_Ik(model, dt):
+    error = np.zeros(model._I.size)
+    h_Ik_next = model.h_Ik
+    h_Ik_prev = model.states['h_Ik']
+    _I_internal = (~model._I_start) & (~model._I_end)
+    error += model._E_Ik * h_Ik_next
+    error -= model._D_Ik
+    error[model._I_start] -= model._Q_uk
+    error[model._I_start] += model._Q_ik[model._i_1k]
+    error[model._I_end] -= model._Q_ik[model._i_nk]
+    error[model._I_end] += model._Q_dk
+    error[_I_internal] -= model._Q_ik[model.backward_I_i[_I_internal]]
+    error[_I_internal] += model._Q_ik[model.forward_I_i[_I_internal]]
+    return error
+
+def momentum_error_ik(model, dt):
+    error = np.zeros(model._i.size)
+    _i_is_start = np.zeros(model._i.size, dtype=np.bool_)
+    _i_is_start[model._i_1k] = True
+    _i_is_end = np.zeros(model._i.size, dtype=np.bool_)
+    _i_is_end[model._i_nk] = True
+    _i_is_internal = (~_i_is_start) & (~_i_is_end)
+    _im1 = (model._i - 1)[_i_is_internal]
+    _ip1 = (model._i + 1)[_i_is_internal]
+    g = 9.81
+    Q_ik_next = model.Q_ik
+    Q_ik_prev = model.states['Q_ik']
+    h_Ik_next = model.h_Ik
+    error -= model._P_ik
+    error -= g * model._A_ik * (h_Ik_next[model._Ik] - h_Ik_next[model._Ip1k])
+    error += model._b_ik * Q_ik_next
+    error[_i_is_start] += model._a_ik[_i_is_start] * model.Q_uk
+    error[_i_is_end] += model._c_ik[_i_is_end] * model.Q_dk
+    error[_i_is_internal] += model._a_ik[_i_is_internal] * model.Q_ik[_im1]
+    error[_i_is_internal] += model._c_ik[_i_is_internal] * model.Q_ik[_ip1]
+    return error
+
+def compute_error(model, dt):
+    error_j = continuity_error_j(model, dt)
+    error_Ik = continuity_error_Ik(model, dt)
+    error_ik = momentum_error_ik(model, dt)
+    error = np.concatenate([error_j, error_Ik, error_ik])
+    return error
+
+def volume_j(model):
+    # Import instance variables
+    _functional = model._functional              # Superlinks with functional area curves
+    _tabular = model._tabular                    # Superlinks with tabular area curves
+    _storage_a = model._storage_a                # Coefficient of functional storage curve
+    _storage_b = model._storage_b                # Exponent of functional storage curve
+    _storage_c = model._storage_c                # Constant of functional storage curve
+    H_j = model.H_j                              # Head at superjunction j
+    _z_inv_j = model._z_inv_j                    # Invert elevation at superjunction j
+    min_depth = model.min_depth                  # Minimum depth allowed at superjunctions/nodes
+    _storage_hs = model._storage_hs
+    _storage_As = model._storage_As
+    _storage_Vs = model._storage_Vs
+    _storage_inds = model._storage_inds
+    _storage_lens = model._storage_lens
+    _storage_js = model._storage_js
+    _storage_codes = model._storage_codes
+    # Compute storage areas
+    _V_sj = np.zeros(model.M, dtype=np.float64)
+    _h_j = np.maximum(H_j - _z_inv_j, min_depth)
+    numba_compute_functional_storage_volumes(_h_j, _V_sj, _storage_a, _storage_b,
+                                                _storage_c, _functional)
+    if _tabular.any():
+        numba_compute_tabular_storage_volumes(_h_j, _V_sj, _storage_hs, _storage_As,
+                                                _storage_Vs, _storage_js, _storage_codes,
+                                                _storage_inds, _storage_lens)
+    return _V_sj
+
+def volume_uk(model):
+    V_uk = model._A_uk * model._dx_uk
+    return V_uk
+
+def volume_dk(model):
+    V_dk = model._A_dk * model._dx_dk
+    return V_dk
+
+def volume_Ik(model):
+    V_Ik = model._A_SIk * model.h_Ik
+    return V_Ik
+
+def volume_ik(model):
+    V_ik = model._A_ik * model._dx_ik
+    return V_ik
+
+def total_volume(model):
+    V_sj = volume_j(model)
+    V_uk = volume_uk(model)
+    V_dk = volume_dk(model)
+    V_o = np.zeros(model.n_o, dtype=np.float64)
+    V_w = np.zeros(model.n_w, dtype=np.float64)
+    V_p = np.zeros(model.n_p, dtype=np.float64)
+    V_Ik = volume_Ik(model)
+    V_ik = volume_ik(model)
+    V_total = np.concatenate([V_sj, V_uk, V_dk, V_o, V_w, V_p, V_Ik, V_ik])
+    return V_total

pipedream_solver/nsuperlink.py

@@ -1558,7 +1558,8 @@ def solve_orifice_flows(self, dt, u=None):
         # TODO: Move this inside numba function
         upstream_ctrl = (H_j[_J_uo] > H_j[_J_do])
         _Qo_max = np.where(upstream_ctrl, _V_sj[_J_uo], _V_sj[_J_do]) / dt
-        _Qo_next = np.sign(_Qo_next) * np.minimum(np.abs(_Qo_next), _Qo_max)
+        # TODO: Check if flow limiter is causing issues
+        #_Qo_next = np.sign(_Qo_next) * np.minimum(np.abs(_Qo_next), _Qo_max)
         # Export instance variables
         self._Qo = _Qo_next
 

pipedream_solver/superlink.py

@@ -8,6 +8,7 @@
 import pipedream_solver.geometry
 import pipedream_solver.storage
 import pipedream_solver.visualization
+from pipedream_solver.callbacks import BaseCallback
 
 class SuperLink():
     """
@@ -293,6 +294,7 @@ def __init__(self, superlinks, superjunctions,
         orifices = copy.deepcopy(orifices)
         weirs = copy.deepcopy(weirs)
         pumps = copy.deepcopy(pumps)
+        self.callbacks = {}
         # TODO: Ensure index is int
         # TODO: This needs to be done for orifices/weirs/pumps as well
         # Ensure nominal direction of superlinks is correct
@@ -420,6 +422,7 @@ def __init__(self, superlinks, superjunctions,
         # Dimensions
         self.nk = np.bincount(self._ki)
         # Create forward and backward indexers
+        # TODO: Check if these are correct for single link
         self.forward_I_I = np.copy(self._I)
         self.forward_I_I[self._Ik] = self._Ip1k
         self.backward_I_I = np.copy(self._I)
@@ -3934,6 +3937,8 @@ def save_state(self):
         self.states['A_dk'] = np.copy(self.A_dk)
         self.states['A_sj'] = np.copy(self.A_sj)
         self.states['V_j'] = np.copy(self.V_j)
+        for _, callback in self.callbacks.items():
+            callback.__on_save_state__()
 
     def load_state(self, states={}, exclude_states=set(),
                    compute_hydraulic_geometries=True):
@@ -3959,6 +3964,8 @@ def load_state(self, states={}, exclude_states=set(),
             self.downstream_hydraulic_geometry()
             self.compute_storage_areas()
             self.node_velocities()
+        for _, callback in self.callbacks.items():
+            callback.__on_load_state__()
 
     def spinup(self, n_steps=100, dt=10, Q_in=None, Q_0Ik=None, reposition_junctions=False,
                reset_counters=True, **kwargs):
@@ -4020,12 +4027,21 @@ def plot_network_3d(self, ax=None, superjunction_signal=None, junction_signal=No
                                  weir_kwargs=weir_kwargs,
                                  pump_kwargs=pump_kwargs))
 
+    def bind_callback(self, callback, key):
+        assert isinstance(callback, BaseCallback)
+        self.callbacks[key] = callback
+
+    def unbind_callback(self, key):
+        return self.callbacks.pop(key)
+
     def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
              first_time=False, implicit=True, banded=False, first_iter=True):
         if first_iter:
             self.save_state()
         if dt is None:
             dt = self._dt
+        else:
+            self._dt = dt
         self._H_bc = H_bc
         self._Q_in = Q_in
         self._Q_0Ik = Q_0Ik
@@ -4126,6 +4142,8 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d
         implicit : bool
             (Deprecated)
         """
+        for _, callback in self.callbacks.items():
+            callback.__on_step_start__()
         if banded is None:
             banded = self.banded
         self._setup_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
@@ -4157,3 +4175,7 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d
                         break
                     H_j_next = np.copy(self.H_j)
         self.iter_elapsed = iter_elapsed
+        for _, callback in self.callbacks.items():
+            callback.__on_step_end__(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                                     first_time=first_time, implicit=implicit, banded=banded,
+                                     first_iter=False)