def get_default_moment_transform(lattice):
if lattice.stencil == D1Q3:
return D1Q3Transform(lattice)
if lattice.stencil == D2Q9:
return D2Q9Lallemand(lattice)
else:
raise LettuceException(
f"No default moment transform for lattice {lattice}.")
def __init__(self, lattice, tau):
self.lattice = lattice
assert lattice.Q == 9, LettuceException("KBC2D only realized for D2Q9")
self.tau = tau
self.beta = 1. / (2 * tau)
# Build a matrix that contains the indices
self.M = torch.zeros([3, 3, 9],
device=lattice.device,
dtype=lattice.dtype)
for i in range(3):
for j in range(3):
self.M[i, j] = lattice.e[:, 0]**i * lattice.e[:, 1]**j
def einsum(self, equation, fields, **kwargs):
"""Einstein summation on local fields."""
input, output = equation.split("->")
inputs = input.split(",")
for i, inp in enumerate(inputs):
if len(inp) == len(fields[i].shape):
pass
elif len(inp) == len(fields[i].shape) - self.D:
inputs[i] += "..."
if not output.endswith("..."):
output += "..."
else:
raise LettuceException("Bad dimension.")
equation = ",".join(inputs) + "->" + output
return torch.einsum(equation, fields, **kwargs)