Save ram ii (#67)

lettuce/equilibrium.py

@@ -1,6 +1,7 @@
 import torch
 
-__all__ = ["Equilibrium", "QuadraticEquilibrium", "IncompressibleQuadraticEquilibrium"]
+__all__ = ["Equilibrium", "QuadraticEquilibrium", "IncompressibleQuadraticEquilibrium",
+           "QuadraticEquilibrium_LessMemory"]
 
 
 class Equilibrium():
@@ -12,15 +13,31 @@ def __init__(self, lattice):
         self.lattice = lattice
 
     def __call__(self, rho, u, *args):
+        exu = torch.tensordot(self.lattice.e, u, dims=1)
+        uxu = self.lattice.einsum("d,d->", [u, u])
         feq = self.lattice.einsum(
             "q,q->q",
             [self.lattice.w,
-             (rho * ((2 * torch.tensordot(self.lattice.e, u, dims=1) - self.lattice.einsum("d,d->", [u, u])) /
-                     (2 * self.lattice.cs ** 2) + 0.5 * (torch.tensordot(self.lattice.e, u, dims=1) /
-                                                         (self.lattice.cs ** 2)) ** 2 + 1))]
+             rho * ((2 * exu - uxu) / (2 * self.lattice.cs ** 2) + 0.5 * (exu / (self.lattice.cs ** 2)) ** 2 + 1)]
         )
         return feq
 
+class QuadraticEquilibrium_LessMemory(QuadraticEquilibrium):
+    """does the same as the normal equilibrium, how ever it uses somewhere around 20% less RAM,
+    but runs about 2% slower on GPU and 11% on CPU
+
+    Use this by setting
+    lattice.equilibrium = QuadraticEquilibrium_LessMemory(lattice)
+    before starting your simulation
+    """
+    def __call__(self, rho, u, *args):
+        return self.lattice.einsum(
+            "q,q->q",
+            [self.lattice.w,
+             (rho * ((2 * torch.tensordot(self.lattice.e, u, dims=1) - self.lattice.einsum("d,d->", [u, u])) / (2 * self.lattice.cs ** 2) + 0.5 * (
+                         torch.tensordot(self.lattice.e, u, dims=1) / (self.lattice.cs ** 2)) ** 2 + 1))]
+        )
+
 
 class IncompressibleQuadraticEquilibrium(Equilibrium):
     def __init__(self, lattice, rho0=1.0):