Quality pass: vectorize HX, add HX Jacobian, extract FlashDrum VLE helper, clean up CSTR

Author: milanofthe

src/pathsim_chem/process/cstr.py

@@ -108,9 +108,6 @@ def __init__(self, V=1.0, F=0.1, k0=1e6, Ea=50000.0, n=1.0,
         self.Cp = Cp
         self.UA = UA
 
-        # derived
-        self.tau = V / F
-
         # rhs of CSTR ode system
         def _fn_d(x, u, t):
             C_A, T = x
@@ -119,35 +116,32 @@ def _fn_d(x, u, t):
             tau = self.V / self.F
             k = self.k0 * np.exp(-self.Ea / (R_GAS * T))
             r = k * C_A**self.n
+            rcp = (-self.dH_rxn) / (self.rho * self.Cp)
+            ua_term = self.UA / (self.rho * self.Cp * self.V)
 
             dC_A = (C_A_in - C_A) / tau - r
-            dT = ((T_in - T) / tau
-                  + (-self.dH_rxn) / (self.rho * self.Cp) * r
-                  + self.UA / (self.rho * self.Cp * self.V) * (T_c - T))
+            dT = (T_in - T) / tau + rcp * r + ua_term * (T_c - T)
 
             return np.array([dC_A, dT])
 
         # jacobian of rhs wrt state [C_A, T]
         def _jc_d(x, u, t):
             C_A, T = x
+
             tau = self.V / self.F
             k = self.k0 * np.exp(-self.Ea / (R_GAS * T))
-
-            # partial derivatives
             dk_dT = k * self.Ea / (R_GAS * T**2)
 
-            # dr/dC_A and dr/dT
             dr_dCA = k * self.n * C_A**(self.n - 1) if C_A > 0 else 0.0
             dr_dT = dk_dT * C_A**self.n
 
             rcp = (-self.dH_rxn) / (self.rho * self.Cp)
             ua_term = self.UA / (self.rho * self.Cp * self.V)
 
-            J = np.array([
-                [-1.0/tau - dr_dCA,         -dr_dT],
+            return np.array([
+                [-1.0/tau - dr_dCA,           -dr_dT],
                 [rcp * dr_dCA,  -1.0/tau + rcp * dr_dT - ua_term]
             ])
-            return J
 
         # output function: well-mixed => outlet = state
         def _fn_a(x, u, t):

src/pathsim_chem/process/flash_drum.py

@@ -35,7 +35,11 @@ class FlashDrum(DynamicalSystem):
 
         \\sum_i \\frac{z_i (K_i - 1)}{1 + \\beta (K_i - 1)} = 0
 
-    For a binary system this has an analytical solution.
+    For a binary system this has the analytical solution:
+
+    .. math::
+
+        \\beta = -\\frac{z_1(K_1 - 1) + z_2(K_2 - 1)}{(K_1 - 1)(K_2 - 1)}
 
     Dynamic States
     ---------------
@@ -97,108 +101,53 @@ def __init__(self, holdup=100.0,
         else:
             x0 = np.array([holdup / 2.0, holdup / 2.0])
 
-        # rhs of flash drum ode
-        def _fn_d(x, u, t):
-            F_in, z_1, T, P = u
-
-            z = np.array([z_1, 1.0 - z_1])
-            N_total = x[0] + x[1]
-
-            # liquid composition from holdup
-            if N_total > 0:
-                x_liq = x / N_total
-            else:
-                x_liq = np.array([0.5, 0.5])
-
-            # K-values from Antoine + Raoult
+        # shared VLE computation
+        def _solve_vle(z, T, P):
+            """Solve binary Rachford-Rice analytically, return (beta, y, x)."""
             P_sat = np.exp(self.antoine_A - self.antoine_B / (T + self.antoine_C))
             K = P_sat / P
 
-            # Rachford-Rice for binary: analytical solution
-            # beta = (z1*(K1-1) + z2*(K2-1)) ... but simpler:
-            # beta = (1 - sum(z/K)) / (sum(z*K) - 1) ... use standard form
-            # For binary: beta = (z1 - 1/K1_eff) / (1 - 1/K1_eff) ...
-            # Actually use: beta from f(beta) = sum zi(Ki-1)/(1+beta(Ki-1)) = 0
-            # For binary with z = x_liq (liquid feed to flash):
-            denom1 = K[0] - 1.0
-            denom2 = K[1] - 1.0
-
-            if abs(denom1) < 1e-12 and abs(denom2) < 1e-12:
-                # no separation
+            d1 = K[0] - 1.0
+            d2 = K[1] - 1.0
+
+            if abs(d1) < 1e-12 and abs(d2) < 1e-12:
                 beta = 0.0
             else:
-                # Solve Rachford-Rice for feed composition z
-                # f(beta) = z[0]*(K[0]-1)/(1+beta*(K[0]-1)) + z[1]*(K[1]-1)/(1+beta*(K[1]-1)) = 0
-                # For binary: beta = -z[0]*(K[0]-1)/((K[0]-1)*(K[1]-1)) ...
-                # Direct: beta = (z[0]*(K[0]-1) + z[1]*(K[1]-1)) isn't right
-                # Newton or analytical: for 2 components
-                # z1(K1-1)/(1+b(K1-1)) + z2(K2-1)/(1+b(K2-1)) = 0
-                # z1(K1-1)(1+b(K2-1)) + z2(K2-1)(1+b(K1-1)) = 0
-                # z1(K1-1) + z1(K1-1)*b*(K2-1) + z2(K2-1) + z2(K2-1)*b*(K1-1) = 0
-                # [z1(K1-1) + z2(K2-1)] + b*(K1-1)*(K2-1)*[z1+z2] = 0
-                # b = -[z1(K1-1) + z2(K2-1)] / [(K1-1)*(K2-1)]
-                num = z[0] * denom1 + z[1] * denom2
-                den = denom1 * denom2
-                if abs(den) < 1e-12:
-                    beta = 0.0
-                else:
-                    beta = -num / den
-
-            # clamp beta to [0, 1]
+                den = d1 * d2
+                beta = -(z[0] * d1 + z[1] * d2) / den if abs(den) > 1e-12 else 0.0
+
             beta = np.clip(beta, 0.0, 1.0)
 
-            # vapor and liquid compositions
             y = z * K / (1.0 + beta * (K - 1.0))
             x_eq = z / (1.0 + beta * (K - 1.0))
 
-            # normalize for safety
-            y_sum = y.sum()
-            x_sum = x_eq.sum()
-            if y_sum > 0:
-                y = y / y_sum
-            if x_sum > 0:
-                x_eq = x_eq / x_sum
-
-            # flow rates
-            V_rate = beta * F_in
-            L_rate = (1.0 - beta) * F_in
-
-            # holdup dynamics
-            dN = F_in * z - V_rate * y - L_rate * x_eq
+            # normalize for numerical safety
+            y_s, x_s = y.sum(), x_eq.sum()
+            if y_s > 0:
+                y = y / y_s
+            if x_s > 0:
+                x_eq = x_eq / x_s
 
-            return dN
+            return beta, y, x_eq
 
-        # algebraic output: compute VLE and return rates + compositions
-        def _fn_a(x, u, t):
+        # rhs of flash drum ode
+        def _fn_d(x, u, t):
             F_in, z_1, T, P = u
-
             z = np.array([z_1, 1.0 - z_1])
 
-            # K-values
-            P_sat = np.exp(self.antoine_A - self.antoine_B / (T + self.antoine_C))
-            K = P_sat / P
-
-            denom1 = K[0] - 1.0
-            denom2 = K[1] - 1.0
+            beta, y, x_eq = _solve_vle(z, T, P)
 
-            if abs(denom1) < 1e-12 and abs(denom2) < 1e-12:
-                beta = 0.0
-            else:
-                num = z[0] * denom1 + z[1] * denom2
-                den = denom1 * denom2
-                beta = -num / den if abs(den) > 1e-12 else 0.0
+            V_rate = beta * F_in
+            L_rate = (1.0 - beta) * F_in
 
-            beta = np.clip(beta, 0.0, 1.0)
+            return F_in * z - V_rate * y - L_rate * x_eq
 
-            y = z * K / (1.0 + beta * (K - 1.0))
-            x_eq = z / (1.0 + beta * (K - 1.0))
+        # algebraic output
+        def _fn_a(x, u, t):
+            F_in, z_1, T, P = u
+            z = np.array([z_1, 1.0 - z_1])
 
-            y_sum = y.sum()
-            x_sum = x_eq.sum()
-            if y_sum > 0:
-                y = y / y_sum
-            if x_sum > 0:
-                x_eq = x_eq / x_sum
+            beta, y, x_eq = _solve_vle(z, T, P)
 
             V_rate = beta * F_in
             L_rate = (1.0 - beta) * F_in

src/pathsim_chem/process/heat_exchanger.py

@@ -105,53 +105,90 @@ def __init__(self, N_cells=5, F_h=0.1, F_c=0.1, V_h=0.5, V_c=0.5,
         self.rho_c = rho_c
         self.Cp_c = Cp_c
 
-        # per-cell quantities
         N = self.N_cells
-        V_cell_h = V_h / N
-        V_cell_c = V_c / N
-        UA_cell = UA / N
 
         # initial state: interleaved [T_h1, T_c1, T_h2, T_c2, ...]
         x0 = np.empty(2 * N)
-        x0[0::2] = T_h0  # hot side
-        x0[1::2] = T_c0  # cold side
+        x0[0::2] = T_h0
+        x0[1::2] = T_c0
 
-        # rhs of heat exchanger ode
+        # rhs of heat exchanger ode (vectorized)
         def _fn_d(x, u, t):
             T_h_in, T_c_in = u
-            dx = np.zeros(2 * N)
+            N = self.N_cells
 
-            _V_cell_h = self.V_h / self.N_cells
-            _V_cell_c = self.V_c / self.N_cells
-            _UA_cell = self.UA / self.N_cells
+            V_cell_h = self.V_h / N
+            V_cell_c = self.V_c / N
+            UA_cell = self.UA / N
 
-            alpha_h = _UA_cell / (self.rho_h * self.Cp_h * _V_cell_h)
-            alpha_c = _UA_cell / (self.rho_c * self.Cp_c * _V_cell_c)
-            flow_h = self.F_h / _V_cell_h
-            flow_c = self.F_c / _V_cell_c
+            fh = self.F_h / V_cell_h
+            fc = self.F_c / V_cell_c
+            ah = UA_cell / (self.rho_h * self.Cp_h * V_cell_h)
+            ac = UA_cell / (self.rho_c * self.Cp_c * V_cell_c)
 
-            for i in range(self.N_cells):
-                T_h_i = x[2*i]
-                T_c_i = x[2*i + 1]
+            T_h = x[0::2]
+            T_c = x[1::2]
 
-                # hot side: upstream is i-1, or inlet
-                T_h_prev = x[2*(i-1)] if i > 0 else T_h_in
-                # cold side: upstream (counter-current) is i+1, or inlet
-                T_c_next = x[2*(i+1) + 1] if i < self.N_cells - 1 else T_c_in
+            # upstream temperatures with boundary conditions
+            T_h_prev = np.empty(N)
+            T_h_prev[0] = T_h_in
+            T_h_prev[1:] = T_h[:-1]
 
-                dx[2*i]     = flow_h * (T_h_prev - T_h_i) - alpha_h * (T_h_i - T_c_i)
-                dx[2*i + 1] = flow_c * (T_c_next - T_c_i) + alpha_c * (T_h_i - T_c_i)
+            T_c_next = np.empty(N)
+            T_c_next[-1] = T_c_in
+            T_c_next[:-1] = T_c[1:]
+
+            dx = np.empty(2 * N)
+            dx[0::2] = fh * (T_h_prev - T_h) - ah * (T_h - T_c)
+            dx[1::2] = fc * (T_c_next - T_c) + ac * (T_h - T_c)
 
             return dx
 
+        # analytical jacobian (block-tridiagonal structure)
+        def _jc_d(x, u, t):
+            N = self.N_cells
+
+            V_cell_h = self.V_h / N
+            V_cell_c = self.V_c / N
+            UA_cell = self.UA / N
+
+            fh = self.F_h / V_cell_h
+            fc = self.F_c / V_cell_c
+            ah = UA_cell / (self.rho_h * self.Cp_h * V_cell_h)
+            ac = UA_cell / (self.rho_c * self.Cp_c * V_cell_c)
+
+            dim = 2 * N
+            J = np.zeros((dim, dim))
+
+            for i in range(N):
+                hi = 2 * i      # hot index
+                ci = 2 * i + 1  # cold index
+
+                # dT_h_i/dT_h_i, dT_h_i/dT_c_i
+                J[hi, hi] = -fh - ah
+                J[hi, ci] = ah
+
+                # dT_c_i/dT_h_i, dT_c_i/dT_c_i
+                J[ci, hi] = ac
+                J[ci, ci] = -fc - ac
+
+                # dT_h_i/dT_h_{i-1} (hot upstream coupling)
+                if i > 0:
+                    J[hi, 2*(i-1)] = fh
+
+                # dT_c_i/dT_c_{i+1} (cold upstream coupling, counter-current)
+                if i < N - 1:
+                    J[ci, 2*(i+1) + 1] = fc
+
+            return J
+
         # output: hot outlet = last cell hot, cold outlet = first cell cold
         def _fn_a(x, u, t):
-            T_h_out = x[2*(self.N_cells - 1)]    # last hot cell
-            T_c_out = x[1]                         # first cold cell
-            return np.array([T_h_out, T_c_out])
+            return np.array([x[2*(self.N_cells - 1)], x[1]])
 
         super().__init__(
             func_dyn=_fn_d,
+            jac_dyn=_jc_d,
             func_alg=_fn_a,
             initial_value=x0,
         )

tests/process/test_cstr.py

@@ -32,7 +32,6 @@ def test_init_default(self):
         self.assertEqual(C.rho, 1000.0)
         self.assertEqual(C.Cp, 4184.0)
         self.assertEqual(C.UA, 500.0)
-        self.assertAlmostEqual(C.tau, 10.0)
 
     def test_init_custom(self):
         """Test custom initialization."""
@@ -41,7 +40,6 @@ def test_init_custom(self):
                  C_A0=1.5, T0=350.0)
         self.assertEqual(C.V, 2.0)
         self.assertEqual(C.F, 0.5)
-        self.assertAlmostEqual(C.tau, 4.0)
 
         C.set_solver(EUF, parent=None)
         self.assertTrue(np.allclose(C.engine.initial_value, [1.5, 350.0]))