def test_initialization(dtype_device, use_jacobi):
dtype, device = dtype_device
lattice = Lattice(D2Q9, device, dtype)
flow = TaylorGreenVortex2D(resolution=24,
reynolds_number=10,
mach_number=0.05,
lattice=lattice)
collision = BGKCollision(lattice, tau=flow.units.relaxation_parameter_lu)
streaming = StandardStreaming(lattice)
simulation = Simulation(flow=flow,
lattice=lattice,
collision=collision,
streaming=streaming)
# set initial pressure to 0 everywhere
p, u = flow.initial_solution(flow.grid)
u0 = lattice.convert_to_tensor(flow.units.convert_velocity_to_lu(u))
rho0 = lattice.convert_to_tensor(np.ones_like(u0[0, ...].cpu()))
simulation.f = lattice.equilibrium(rho0, u0)
if use_jacobi:
simulation.initialize_pressure(1000, 1e-6)
num_iterations = 0
else:
num_iterations = simulation.initialize(500, 1e-3)
piter = lattice.convert_to_numpy(
flow.units.convert_density_lu_to_pressure_pu(lattice.rho(
simulation.f)))
# assert that pressure is converged up to 0.05 (max p
assert piter == pytest.approx(p, rel=0.0, abs=5e-2)
assert num_iterations < 500
enstrophy_f = ObservableReporter(Enstrophy(lattice, flow_fine),
interval=postplotinterval * 2,
out=None)
energy_f = ObservableReporter(IncompressibleKineticEnergy(lattice, flow_fine),
interval=postplotinterval * 2,
out=None)
spectrum_f = ObservableReporter(EnergySpectrum(lattice, flow_fine),
interval=postplotinterval * 2,
out=None)
f_reporter = FReporter(2 * data_interval)
simulation.reporters.append(f_reporter)
simulation.reporters.append(enstrophy_f)
simulation.reporters.append(energy_f)
simulation.reporters.append(spectrum_f)
simulation.initialize_pressure()
simulation.initialize_f_neq()
_ = simulation.step(2 * n_steps)
########################################################################################################################
##### Make Data Pairs
########################################################################################################################
print('Make Data Pairs and shuffle training data')
data = torch.stack(f_reporter.fs).cpu()
f_reporter.fs = []
inputs = data[0:-1]
outputs = data[1:]
f_init = inputs[nr_init].clone()
f_output = outputs[-1].clone()
