Fixed pressure computation

MagriLab · Oct 24, 2023 · 4ef889a · 4ef889a
1 parent ad3c724
commit 4ef889a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/kolsol/numpy/solver.py b/kolsol/numpy/solver.py
@@ -196,7 +196,7 @@ def pressure(self, u_hat: np.ndarray) -> np.ndarray:
         f_hat = oe.contract('...t, ut... -> ...u', -self.nabla, aapt)
 
         with np.errstate(divide='ignore', invalid='ignore'):
-            p_hat = oe.contract('...u, ...u -> ...', -self.nabla, f_hat) / self.kk
+            p_hat = oe.contract('...u, ...u -> ...', -self.nabla, f_hat + self.f) / self.kk
             p_hat[tuple([...]) + tuple(self.nk for _ in range(self.ndim))] = 0.0
 
         return p_hat

diff --git a/kolsol/torch/solver.py b/kolsol/torch/solver.py
@@ -240,7 +240,7 @@ def pressure(self, u_hat: torch.Tensor) -> torch.Tensor:
         aapt = torch.stack(uij_aapt, dim=0)
         f_hat = oe.contract('...t, ut... -> ...u', -self.nabla, aapt)
 
-        p_hat = oe.contract('...u, ...u -> ...', -self.nabla, f_hat) / self.kk_div
+        p_hat = oe.contract('...u, ...u -> ...', -self.nabla, f_hat + self.f) / self.kk_div
         p_hat[tuple([...]) + tuple(self.nk for _ in range(self.ndim))] = 0.0
 
         return p_hat