"""
Observables.
Each observable is defined as a callable class.
The `__call__` function takes f as an argument and returns a torch tensor.
"""
import torch
import numpy as np
from lettuce.util import torch_gradient
from packaging import version
__all__ = ["Observable", "MaximumVelocity", "IncompressibleKineticEnergy", "Enstrophy", "EnergySpectrum"]
[docs]
class Observable:
def __init__(self, lattice, flow):
self.lattice = lattice
self.flow = flow
def __call__(self, f):
raise NotImplementedError
[docs]
class MaximumVelocity(Observable):
"""Maximum velocitiy"""
def __call__(self, f):
u = self.lattice.u(f)
return self.flow.units.convert_velocity_to_pu(torch.norm(u, dim=0).max())
[docs]
class IncompressibleKineticEnergy(Observable):
"""Total kinetic energy of an incompressible flow."""
def __call__(self, f):
dx = self.flow.units.convert_length_to_pu(1.0)
kinE = self.flow.units.convert_incompressible_energy_to_pu(torch.sum(self.lattice.incompressible_energy(f)))
kinE *= dx ** self.lattice.D
return kinE
[docs]
class Enstrophy(Observable):
"""The integral of the vorticity
Notes
-----
The function only works for periodic domains
"""
def __call__(self, f):
u0 = self.flow.units.convert_velocity_to_pu(self.lattice.u(f)[0])
u1 = self.flow.units.convert_velocity_to_pu(self.lattice.u(f)[1])
dx = self.flow.units.convert_length_to_pu(1.0)
grad_u0 = torch_gradient(u0, dx=dx, order=6)
grad_u1 = torch_gradient(u1, dx=dx, order=6)
vorticity = torch.sum((grad_u0[1] - grad_u1[0]) * (grad_u0[1] - grad_u1[0]))
if self.lattice.D == 3:
u2 = self.flow.units.convert_velocity_to_pu(self.lattice.u(f)[2])
grad_u2 = torch_gradient(u2, dx=dx, order=6)
vorticity += torch.sum(
(grad_u2[1] - grad_u1[2]) * (grad_u2[1] - grad_u1[2])
+ ((grad_u0[2] - grad_u2[0]) * (grad_u0[2] - grad_u2[0]))
)
return vorticity * dx ** self.lattice.D
[docs]
class EnergySpectrum(Observable):
"""The kinetic energy spectrum"""
def __init__(self, lattice, flow):
super(EnergySpectrum, self).__init__(lattice, flow)
self.dx = self.flow.units.convert_length_to_pu(1.0)
self.dimensions = self.flow.grid[0].shape
frequencies = [self.lattice.convert_to_tensor(np.fft.fftfreq(dim, d=1 / dim)) for dim in self.dimensions]
wavenumbers = torch.stack(torch.meshgrid(*frequencies))
wavenorms = torch.norm(wavenumbers, dim=0)
if self.lattice.D == 3:
self.norm = self.dimensions[0] * np.sqrt(2 * np.pi) / self.dx ** 2
else:
self.norm = self.dimensions[0] / self.dx
self.wavenumbers = torch.arange(int(torch.max(wavenorms)))
self.wavemask = (
(wavenorms[..., None] > self.wavenumbers.to(dtype=lattice.dtype, device=lattice.device) - 0.5) &
(wavenorms[..., None] <= self.wavenumbers.to(dtype=lattice.dtype, device=lattice.device) + 0.5)
)
def __call__(self, f):
u = self.lattice.u(f)
return self.spectrum_from_u(u)
[docs]
def spectrum_from_u(self, u):
u = self.flow.units.convert_velocity_to_pu(u)
ekin = self._ekin_spectrum(u)
ek = ekin[..., None] * self.wavemask.to(dtype=self.lattice.dtype)
ek = ek.sum(torch.arange(self.lattice.D).tolist())
return ek
def _ekin_spectrum(self, u):
"""distinguish between different torch versions"""
torch_ge_18 = (version.parse(torch.__version__) >= version.parse("1.8.0"))
if torch_ge_18:
return self._ekin_spectrum_torch_ge_18(u)
else:
return self._ekin_spectrum_torch_lt_18(u)
def _ekin_spectrum_torch_lt_18(self, u):
zeros = torch.zeros(self.dimensions, dtype=self.lattice.dtype, device=self.lattice.device)[..., None]
uh = (torch.stack([
torch.fft(torch.cat((u[i][..., None], zeros), self.lattice.D),
signal_ndim=self.lattice.D) for i in range(self.lattice.D)]) / self.norm)
ekin = torch.sum(0.5 * (uh[..., 0] ** 2 + uh[..., 1] ** 2), dim=0)
return ekin
def _ekin_spectrum_torch_ge_18(self, u):
uh = (torch.stack([
torch.fft.fftn(u[i], dim=tuple(torch.arange(self.lattice.D))) for i in range(self.lattice.D)
]) / self.norm)
ekin = torch.sum(0.5 * (uh.imag ** 2 + uh.real ** 2), dim=0)
return ekin
class Mass(Observable):
"""Total mass in lattice units.
Parameters
----------
no_mass_mask : torch.Tensor
Boolean mask that defines grid points
which do not count into the total mass (e.g. bounce-back boundaries).
"""
def __init__(self, lattice, flow, no_mass_mask=None):
super(Mass, self).__init__(lattice, flow)
self.mask = no_mass_mask
def __call__(self, f):
mass = f[..., 1:-1, 1:-1].sum()
if self.mask is not None:
mass -= (f * self.mask.to(dtype=torch.float)).sum()
return mass