def sample(value, domain, batch_size=None, name=None):
assert isinstance(domain, Domain)
if isinstance(value, Field):
assert_same_rank(
value.rank, domain.rank,
'rank of value (%s) does not match domain (%s)' %
(value.rank, domain.rank))
if isinstance(value,
CenteredGrid) and value.box == domain.box and np.all(
value.resolution == domain.resolution):
data = value.data
else:
point_field = CenteredGrid.getpoints(domain.box,
domain.resolution)
point_field._batch_size = batch_size
data = value.at(point_field).data
else: # value is constant
if callable(value):
x = CenteredGrid.getpoints(
domain.box, domain.resolution).copied_with(
extrapolation=Material.extrapolation_mode(
domain.boundaries),
name=name)
value = value(x)
return value
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.add(
math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )), value)
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries),
name=name)
def accessible(self, accessible):
if accessible is not None:
assert isinstance(accessible, CenteredGrid)
assert accessible.rank == self.domain.rank
return accessible
else:
return self.domain.centered_grid(1, extrapolation=Material.extrapolation_mode(self.domain.boundaries))
def active(self, active):
extrapolation = _active_extrapolation(
Material.extrapolation_mode(self.domain.boundaries))
if active is not None:
assert isinstance(active, CenteredGrid)
assert active.rank == self.domain.rank
if active.extrapolation != extrapolation:
active = active.copied_with(extrapolation=extrapolation)
return active
else:
return self.domain.centered_grid(1, extrapolation=extrapolation)
def solve(self, divergence, domain, guess, enable_backprop):
assert isinstance(domain, PoissonDomain)
fluid_mask = domain.accessible_tensor(extend=1)
extrapolation = Material.extrapolation_mode(domain.domain.boundaries)
def apply_A(pressure):
pressure = CenteredGrid(pressure, extrapolation=extrapolation)
pressure_padded = pressure.padded([[1, 1]] * pressure.rank)
return _weighted_sliced_laplace_nd(pressure_padded.data,
weights=fluid_mask)
return conjugate_gradient(divergence,
apply_A,
guess,
self.accuracy,
self.max_iterations,
back_prop=enable_backprop)
def solve_pressure_forward(divergence,
fluid_mask,
max_iterations,
guess,
accuracy,
domain,
back_prop=False):
from phi.physics.material import Material
extrapolation = Material.extrapolation_mode(domain.domain.boundaries)
def apply_A(pressure):
pressure = CenteredGrid(pressure, extrapolation=extrapolation)
pressure_padded = pressure.padded([[1, 1]] * pressure.rank)
return _weighted_sliced_laplace_nd(pressure_padded.data,
weights=fluid_mask)
return conjugate_gradient(divergence,
apply_A,
guess,
accuracy,
max_iterations,
back_prop=back_prop)
def sample(value, domain, batch_size=None):
assert isinstance(domain, Domain)
if isinstance(value, Field):
assert_same_rank(
value.rank, domain.rank,
'rank of value (%s) does not match domain (%s)' %
(value.rank, domain.rank))
if isinstance(value,
CenteredGrid) and value.box == domain.box and np.all(
value.resolution == domain.resolution):
data = value.data
else:
data = value.sample_at(
CenteredGrid.getpoints(domain.box, domain.resolution).data)
else: # value is constant
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )) + value
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries))
