Commit 67e0607d authored by Markus Holzer's avatar Markus Holzer
Browse files

Improve getter and setter API

parent 1c008ce4
Pipeline #47585 passed with stages
in 144 minutes and 40 seconds
import functools
from copy import deepcopy
from lbmpy.simplificationfactory import create_simplification_strategy
import functools
from typing import Union, Tuple
from warnings import warn
from pystencils import create_kernel, CreateKernelConfig
from pystencils.field import Field, get_layout_of_array
from pystencils.enums import Target
from lbmpy.advanced_streaming.utility import get_accessor, Timestep
from lbmpy.methods import AbstractLbMethod, LbmCollisionRule
from lbmpy.simplificationfactory import create_simplification_strategy
from lbmpy.stencils import LBStencil
def macroscopic_values_setter(lb_method: AbstractLbMethod, density: Union[Field, int, float],
velocity: Union[Field, Tuple[int], Tuple[float]], pdfs: Field,
streaming_pattern: str = 'pull', previous_timestep: Timestep = Timestep.BOTH,
set_pre_collision_pdfs: bool = False) -> LbmCollisionRule:
"""Assignments to initialize the pdf field with equilibrium
"""
stencil = lb_method.stencil
pdfs_writes = _get_pdfs_field_accesses(pdfs=pdfs, stencil=stencil, streaming_pattern=streaming_pattern,
previous_timestep=previous_timestep,
set_pre_collision_pdfs=set_pre_collision_pdfs)
if not isinstance(density, (int, float)):
if isinstance(density, Field):
density = density.center
else:
warn("In future releases only integer, float or instance of pystencils.Field will be supported to set "
"the density", DeprecationWarning)
if isinstance(velocity, tuple):
if any(isinstance(v, Field.Access) for v in velocity):
warn("Providing velocity as tuple of Field.Access will not be supported in future releases",
DeprecationWarning)
if not all(isinstance(v, (int, float)) for v in velocity):
raise ValueError("If velocity is provided as tuple only tuples of numerical values are supported")
if isinstance(velocity, Field):
velocity = velocity.center_vector
cqc = lb_method.conserved_quantity_computation
inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity, force_substitution=False)
setter_eqs = lb_method.get_equilibrium(conserved_quantity_equations=inp_eqs)
setter_eqs = setter_eqs.new_with_substitutions({sym: pdfs_writes[i]
for i, sym in enumerate(lb_method.post_collision_pdf_symbols)})
return setter_eqs
def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
streaming_pattern='pull', previous_timestep=Timestep.BOTH,
set_pre_collision_pdfs=False):
warn("pdf_initialization_assignments will not be supported in future releases. "
"Please switch to macroscopic_values_setter.", DeprecationWarning)
"""Assignments to initialize the pdf field with equilibrium"""
if isinstance(pdfs, Field):
accessor = get_accessor(streaming_pattern, previous_timestep)
......@@ -38,31 +83,34 @@ def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
return setter_eqs
def macroscopic_values_getter(lb_method, density, velocity, pdfs,
streaming_pattern='pull', previous_timestep=Timestep.BOTH,
use_pre_collision_pdfs=False):
if isinstance(pdfs, Field):
accessor = get_accessor(streaming_pattern, previous_timestep)
if use_pre_collision_pdfs:
field_accesses = accessor.read(pdfs, lb_method.stencil)
else:
field_accesses = accessor.write(pdfs, lb_method.stencil)
elif streaming_pattern == 'pull' and not use_pre_collision_pdfs:
field_accesses = pdfs
else:
raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
+ f"getter assignments for streaming pattern {streaming_pattern}.")
def macroscopic_values_getter(lb_method: AbstractLbMethod,
density: Field = None, velocity: Field = None, pdfs: Field = None,
streaming_pattern: str = 'pull', previous_timestep: Timestep = Timestep.BOTH,
use_pre_collision_pdfs: bool = False) -> LbmCollisionRule:
assert pdfs is not None, "Please provide a Field to read the PDFs from"
stencil = lb_method.stencil
pdfs_writes = _get_pdfs_field_accesses(pdfs=pdfs, stencil=stencil, streaming_pattern=streaming_pattern,
previous_timestep=previous_timestep,
set_pre_collision_pdfs=use_pre_collision_pdfs)
cqc = lb_method.conserved_quantity_computation
assert not (velocity is None and density is None)
output_spec = {}
if velocity is not None:
if not isinstance(velocity, Field):
warn("In future releases only instance of pystencils.Field will be supported to write the velocity ",
DeprecationWarning)
output_spec['velocity'] = velocity
if density is not None:
warn("In future releases only instance of pystencils.Field will be supported to write the density ",
DeprecationWarning)
output_spec['density'] = density
return cqc.output_equations_from_pdfs(field_accesses, output_spec)
macroscopic_values_setter = pdf_initialization_assignments
ac = cqc.output_equations_from_pdfs(pdfs_writes, output_spec)
result = LbmCollisionRule(lb_method=lb_method,
main_assignments=ac.main_assignments, subexpressions=ac.subexpressions,
simplification_hints=ac.simplification_hints)
return result
def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
......@@ -260,3 +308,22 @@ def create_advanced_velocity_setter_collision_rule(method, velocity_field: Field
new_collision_rule = simplification_strategy(method.get_collision_rule(eq_input))
return new_collision_rule
def _get_pdfs_field_accesses(pdfs: Field, stencil: LBStencil,
streaming_pattern: str = 'pull', previous_timestep: Timestep = Timestep.BOTH,
set_pre_collision_pdfs: bool = False):
if isinstance(pdfs, Field):
accessor = get_accessor(streaming_pattern, previous_timestep)
if set_pre_collision_pdfs:
field_accesses = accessor.read(pdfs, stencil)
else:
field_accesses = accessor.write(pdfs, stencil)
elif streaming_pattern == 'pull' and not set_pre_collision_pdfs:
warn("Only providing the pdfs as instance of Field will be supported in future releases", DeprecationWarning)
field_accesses = pdfs
else:
raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
+ f"initialization assignments for streaming pattern {streaming_pattern}.")
return field_accesses
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment