def _generate_examples():
# --- Example 1 ---
ex1 = np.tile(np.linspace(1, 0, 5), [4, 1])
ex1 = math.expand_dims(math.expand_dims(ex1, -1), 0) - math.mean(ex1)
# --- Example 2 ---
ex2 = np.zeros([1, 4, 5, 1])
ex2[0, :, 2, 0] = 1
ex2 -= math.mean(ex2)
# --- Stack examples to batch ---
return math.concat([ex1, ex2], axis=0)
def points(self,
points: Tensor or Number or tuple or list,
values: Tensor or Number = None,
radius: Tensor or float or int or None = None,
extrapolation: math.Extrapolation = math.extrapolation.ZERO,
color: str or Tensor or tuple or list or None = None) -> PointCloud:
"""
Create a `phi.field.PointCloud` from the given `points`.
The created field has no channel dimensions and all points carry the value `1`.
Args:
points: point locations in physical units
values: (optional) values of the particles, defaults to 1.
radius: (optional) size of the particles
extrapolation: (optional) extrapolation to use, defaults to extrapolation.ZERO
color: (optional) color used when plotting the points
Returns:
`phi.field.PointCloud` object
"""
extrapolation = extrapolation if isinstance(extrapolation, math.Extrapolation) else self.boundaries[extrapolation]
if radius is None:
radius = math.mean(self.bounds.size) * 0.005
# --- Parse points: tuple / list ---
if isinstance(points, (tuple, list)):
if len(points) == 0: # no points
points = math.zeros(instance(points=0), channel(vector=1))
elif isinstance(points[0], Number): # single point
points = math.tensor([points], instance('points'), channel('vector'))
else:
points = math.tensor(points, instance('points'), channel('vector'))
elements = Sphere(points, radius)
if values is None:
values = math.tensor(1.)
return PointCloud(elements, values, extrapolation, add_overlapping=False, bounds=self.bounds, color=color)
def respect_boundaries(particles: PointCloud,
domain: Domain,
not_accessible: list,
offset: float = 0.5) -> PointCloud:
"""
Enforces boundary conditions by correcting possible errors of the advection step and shifting particles out of
obstacles or back into the domain.
Args:
particles: PointCloud holding particle positions as elements
domain: Domain for which any particles outside should get shifted inwards
not_accessible: List of Obstacle or Geometry objects where any particles inside should get shifted outwards
offset: Minimum distance between particles and domain boundary / obstacle surface after particles have been shifted.
Returns:
PointCloud where all particles are inside the domain / outside of obstacles.
"""
new_positions = particles.elements.center
for obj in not_accessible:
if isinstance(obj, Obstacle):
obj = obj.geometry
new_positions = obj.push(new_positions, shift_amount=offset)
new_positions = (~domain.bounds).push(new_positions, shift_amount=offset)
return particles.with_(
elements=Sphere(new_positions,
math.mean(particles.bounds.size) * 0.005))
def grid_sample(self, resolution, size, batch_size=1, channels=None):
channels = channels or self.channels or len(size)
shape = (batch_size, ) + tuple(resolution) + (channels, )
rndj = math.to_complex(
self.math.random_normal(shape)) + 1j * math.to_complex(
self.math.random_normal(shape)) # Note: there is no complex32
k = math.fftfreq(
resolution) * resolution / size * self.scale # in physical units
k = math.sum(k**2, axis=-1, keepdims=True)
lowest_frequency = 0.1
weight_mask = 1 / (1 + math.exp(
(lowest_frequency - k) * 1e3)) # High pass filter
# --- Compute 1/k ---
k[(0, ) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k[(0, ) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k**self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array,
axis=tuple(range(1, math.ndims(array))),
keepdims=True)
array -= math.mean(array,
axis=tuple(range(1, math.ndims(array))),
keepdims=True)
array = math.to_float(array)
return array
def log_scalars(self, frame: int, **values: float or math.Tensor):
"""
Adds `values` to the curves by name.
This can be used to log the evolution of scalar quantities or summaries.
The values are stored in a text file within the scene directory.
The curves may also be directly viewed in the user interface.
Args:
frame: step
values: Values and names to append to the curves, must be numbers or `phi.math.Tensor`.
If a curve does not yet exists, a new one is created.
"""
for name, value in values.items():
assert isinstance(name, str)
value = float(math.mean(value).mean)
if name not in self._scalars:
self._scalars[name] = []
if self.scene is not None:
path = self.scene.subpath(f"log_{name}.txt")
self._scalar_streams[name] = open(path, "w")
self._scalars[name].append((frame, value))
if self.scene is not None:
self._scalar_streams[name].write(f"{frame} {value}\n")
self._scalar_streams[name].flush()
def grid_sample(self,
resolution: math.Shape,
size,
shape: math.Shape = None):
shape = (self._shape if shape is None else shape) & resolution
for dim in channel(self._shape):
if dim.item_names[0] is None:
warnings.warn(
f"Please provide item names for Noise dim {dim} using {dim}='x,y,z'",
FutureWarning)
shape &= channel(**{dim.name: resolution.names})
rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(
random_normal(shape)) # Note: there is no complex32
with math.NUMPY:
k = math.fftfreq(resolution) * resolution / math.tensor(
size) * math.tensor(self.scale) # in physical units
k = math.vec_squared(k)
lowest_frequency = 0.1
weight_mask = math.to_float(k > lowest_frequency)
# --- Compute 1/k ---
k._native[(0, ) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k._native[(0, ) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k**self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array, dim=array.shape.non_batch)
array -= math.mean(array, dim=array.shape.non_batch)
array = math.to_float(array)
return array
def normalize(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
global_mean = math.mean(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
global_mean = math.expand_dims(global_mean, -1)
return StaggeredGrid(self.staggered / global_mean)
def mean(field: SampledField) -> Tensor:
"""
Computes the mean value by reducing all spatial / instance dimensions.
Args:
field: `SampledField`
Returns:
`phi.Tensor`
"""
return math.mean(field.values, field.shape.non_channel.non_batch)
def test_grid_sample_gradient_1d(self):
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
grid = math.tensor([0., 1, 2, 3], spatial('x'))
coords = math.tensor([0.5, 1.5], instance('points'))
with math.record_gradients(grid, coords):
sampled = math.grid_sample(grid, coords, extrapolation.ZERO)
loss = math.mean(math.l2_loss(sampled)) / 2
grad_grid, grad_coords = math.gradients(loss, grid, coords)
math.assert_close(grad_grid, math.tensor([0.125, 0.5, 0.375, 0], spatial('x')), msg=backend)
math.assert_close(grad_coords, math.tensor([0.25, 0.75], instance('points')), msg=backend)
def step(self, u, dt=1.0, **dependent_states):
assert isinstance(u, CenteredGrid)
grad = u.gradient()
laplace = u.laplace()
laplace2 = laplace.laplace()
du_dt = -laplace - laplace2 - 0.5 * grad**2
result = u + dt * du_dt
result -= math.mean(result.data,
axis=tuple(
range(1, len(math.staticshape(result.data)))),
keepdims=True)
return result.copied_with(age=u.age + dt, name=u.name)
def _run_higher_order_fft_reconstruction(in_field, set_accuracy, tolerance=20, order=2):
# Higher Order FFT test
mean = math.mean(in_field).data
centered_field = in_field - mean
fft_poisson = math.fourier_poisson(in_field, times=order)
fft_poisson += mean
fft_reconst = math.fourier_laplace(fft_poisson)
for _ in range(order - 1):
fft_reconst = math.fourier_laplace(fft_reconst)
error = (in_field - fft_reconst) / in_field
max_error = np.max(np.abs(error.data))
passed = max_error < tolerance * set_accuracy
print("{:.2g} vs. {:.2g}".format(max_error, tolerance * set_accuracy))
def _test_reconstruction_first_order(in_field, solve_func, laplace_func, set_accuracy, name, first_order_tolerance=2):
# Test Reconstruction
mean = math.mean(in_field).data
centered_field = in_field - mean
ret = solve_func(centered_field)
try:
solved_field, it = ret
except:
solved_field = ret
reconst1 = laplace_func(solved_field) + mean # Reconstruct Input
error = (in_field - reconst1) / in_field
max_error = np.max(np.abs(error.data))
print("{:.2g}/{:.2g}".format(max_error, first_order_tolerance * set_accuracy))
assert max_error < first_order_tolerance * set_accuracy, "{} reconstruction not within set accuracy. {:.2g} vs. {:.2g}".format(name, max_error, first_order_tolerance * set_accuracy)
def plot_solves():
"""
While `plot_solves()` is active, certain performance optimizations and algorithm implementations may be disabled.
"""
from . import math
import pylab
cycle = pylab.rcParams['axes.prop_cycle'].by_key()['color']
with math.SolveTape(record_trajectories=True) as solves:
try:
yield solves
finally:
for i, result in enumerate(solves):
assert isinstance(result, math.SolveInfo)
from phi.math._tensors import disassemble_tree
_, (residual, ) = disassemble_tree(result.residual)
residual_mse = math.mean(math.sqrt(math.sum(residual**2)),
residual.shape.without('trajectory'))
residual_mse_max = math.max(
math.sqrt(math.sum(residual**2)),
residual.shape.without('trajectory'))
# residual_mean = math.mean(math.abs(residual), residual.shape.without('trajectory'))
residual_max = math.max(math.abs(residual),
residual.shape.without('trajectory'))
pylab.plot(residual_mse.numpy(),
label=f"{i}: {result.method}",
color=cycle[i % len(cycle)])
pylab.plot(residual_max.numpy(),
'--',
alpha=0.2,
color=cycle[i % len(cycle)])
pylab.plot(residual_mse_max.numpy(),
alpha=0.2,
color=cycle[i % len(cycle)])
print(
f"Solve {i}: {result.method} ({1000 * result.solve_time:.1f} ms)\n"
f"\t{result.solve}\n"
f"\t{result.msg}\n"
f"\tConverged: {result.converged}\n"
f"\tDiverged: {result.diverged}\n"
f"\tIterations: {result.iterations}\n"
f"\tFunction evaulations: {result.function_evaluations.trajectory[-1]}"
)
pylab.yscale('log')
pylab.ylabel("Residual: MSE / max / individual max")
pylab.xlabel("Iteration")
pylab.title(f"Solve Convergence")
pylab.legend(loc='upper right')
pylab.savefig(f"pressure-solvers-FP32.png")
pylab.show()
def grid_sample(self, resolution, size, batch_size=1, dtype=np.float32):
shape = (batch_size,) + tuple(resolution) + (self.channels,)
rndj = math.randn(shape, dtype) + 1j * math.randn(shape, dtype)
k = math.fftfreq(resolution) * resolution / size * self.scale # in physical units
k = math.sum(k ** 2, axis=-1, keepdims=True)
lowest_frequency = 0.1
weight_mask = 1 / (1 + math.exp((lowest_frequency - k) * 1e3)) # High pass filter
# --- Compute 1/k ---
k[(0,) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k[(0,) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k ** self.smoothness * weight_mask
array = math.real(math.ifft(fft)).astype(dtype)
array /= math.std(array, axis=tuple(range(1, math.ndims(array))), keepdims=True)
array -= math.mean(array, axis=tuple(range(1, math.ndims(array))), keepdims=True)
return array
def _sample(self, geometry: Geometry) -> Tensor:
if geometry == self.bounds:
return math.mean(self._values, self._resolution)
if isinstance(geometry, GeometryStack):
sampled = [self.sample(g) for g in geometry.geometries]
return math.stack(sampled, geometry.stack_dim)
if isinstance(geometry, GridCell):
if self.elements == geometry:
return self.values
elif math.close(self.dx, geometry.size):
fast_resampled = self._shift_resample(geometry.resolution,
geometry.bounds)
if fast_resampled is not NotImplemented:
return fast_resampled
points = geometry.center
local_points = self.box.global_to_local(points) * self.resolution - 0.5
return math.grid_sample(self.values, local_points, self.extrapolation)
def _plot_points(axis, data: PointCloud, **plt_args):
x, y = [d.numpy() for d in data.points.vector.unstack_spatial('x,y')]
color = [d.native() for d in data.color.points.unstack(len(x))]
if isinstance(data.elements, Sphere):
symbol = 'o'
size = data.elements.bounding_radius().numpy() * 1.41
elif isinstance(data.elements, BaseBox):
symbol = 's'
size = math.mean(data.elements.bounding_half_extent(), 'vector').numpy()
elif isinstance(data.elements, Point):
symbol = 'x'
size = 6 / _get_pixels_per_unit(axis.figure, axis)
else:
symbol = '*'
size = data.elements.bounding_radius().numpy()
size_px = size * _get_pixels_per_unit(axis.figure, axis)
axis.scatter(x, y, marker=symbol, color=color, s=size_px ** 2, alpha=0.8)
_annotate_points(axis, data.points, instance(data))
def grid_sample(self, resolution: math.Shape, size, shape: math.Shape = None):
shape = (self._shape if shape is None else shape).combined(resolution)
rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(random_normal(shape)) # Note: there is no complex32
with math.NUMPY_BACKEND:
k = math.fftfreq(resolution) * resolution / size * self.scale # in physical units
k = math.vec_squared(k)
lowest_frequency = 0.1
weight_mask = 1 / (1 + math.exp((lowest_frequency - k) * 1e3)) # High pass filter
# --- Compute 1/k ---
k.native()[(0,) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k.native()[(0,) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k ** self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array, dim=array.shape.non_batch)
array -= math.mean(array, dim=array.shape.non_batch)
array = math.to_float(array)
return array
def poisson_solve(input_field, poisson_domain, solver=None, guess=None, gradient='implicit'):
"""
Solves the Poisson equation Δp = input_field for p.
:param gradient: one of ('implicit', 'autodiff', 'inverse')
If 'autodiff', use the built-in autodiff for backpropagation.
The intermediate results of each loop iteration will be permanently stored if backpropagation is used.
If 'implicit', computes a forward pressure solve in reverse accumulation backpropagation.
This requires less memory but is only accurate if the solution is fully converged.
:param input_field: CenteredGrid
:param poisson_domain: PoissonDomain instance
:param solver: PoissonSolver to use, None for default
:param guess: CenteredGrid with same size and resolution as input_field
:return: p as CenteredGrid, iteration count as int or None if not available
:rtype: CenteredGrid, int
"""
assert isinstance(input_field, CenteredGrid)
if guess is not None:
assert isinstance(guess, CenteredGrid)
assert guess.compatible(input_field)
guess = guess.data
if isinstance(poisson_domain, Domain):
poisson_domain = PoissonDomain(poisson_domain)
if solver is None:
solver = _choose_solver(input_field.resolution, math.choose_backend([input_field.data, poisson_domain.active.data, poisson_domain.accessible.data]))
if not struct.any(Material.open(poisson_domain.domain.boundaries)): # has no open boundary
input_field = input_field - math.mean(input_field.data, axis=tuple(range(1, 1 + input_field.rank)), keepdims=True) # Subtract mean divergence
assert gradient in ('autodiff', 'implicit', 'inverse')
if gradient == 'autodiff':
pressure, iteration = solver.solve(input_field.data, poisson_domain, guess, enable_backprop=True)
else:
if gradient == 'implicit':
def poisson_gradient(_op, grad):
return poisson_solve(CenteredGrid.sample(grad, poisson_domain.domain), poisson_domain, solver, None, gradient=gradient)[0].data
else: # gradient = 'inverse'
def poisson_gradient(_op, grad):
return CenteredGrid.sample(grad, poisson_domain.domain).laplace(physical_units=False).data
pressure, iteration = math.with_custom_gradient(solver.solve, [input_field.data, poisson_domain, guess, False], poisson_gradient, input_index=0, output_index=0, name_base='poisson_solve')
pressure = CenteredGrid(pressure, input_field.box, extrapolation=input_field.extrapolation, name='pressure')
return pressure, iteration
def _test_reconstruction_second_order(in_field, solve_func, laplace_func, set_accuracy, name, second_order_tolerance=20):
# Calculate 1st order
mean = math.mean(in_field).data
centered_field = in_field - mean
ret = solve_func(centered_field)
try:
solved_field, it = ret
except:
solved_field = ret
# Calculate 2nd order
ret2 = solve_func(solved_field)
try:
solved_field2, it = ret2
except:
solved_field2 = ret2
reconst2 = laplace_func(laplace_func(solved_field2)) + mean
error2 = (in_field - reconst2) / in_field
max_error2 = np.max(np.abs(error2.data))
passed = max_error2 < second_order_tolerance*set_accuracy
print("{}^2 reconstruction {}within set accuracy: {:.2g} vs. {:.2g}".format(name, 'NOT ' if not passed else '', max_error2, second_order_tolerance * set_accuracy))
def divergence_free(velocity,
domain=None,
obstacles=(),
pressure_solver=None,
return_info=False):
"""
Projects the given velocity field by solving for and subtracting the pressure.
:param return_info: if True, returns a dict holding information about the solve as a second object
:param velocity: StaggeredGrid
:param domain: Domain matching the velocity field, used for boundary conditions
:param obstacles: list of Obstacles
:param pressure_solver: PressureSolver. Uses default solver if none provided.
:return: divergence-free velocity as StaggeredGrid
"""
assert isinstance(velocity, StaggeredGrid)
# --- Set up FluidDomain ---
if domain is None:
domain = Domain(velocity.resolution, OPEN)
obstacle_mask = mask(union([obstacle.geometry for obstacle in obstacles]),
antialias=False)
if obstacle_mask is not None:
obstacle_grid = obstacle_mask.at(
velocity.center_points).copied_with(extrapolation='constant')
active_mask = 1 - obstacle_grid
else:
active_mask = math.ones(
domain.centered_shape(name='active', extrapolation='constant'))
accessible_mask = active_mask.copied_with(
extrapolation=Material.accessible_extrapolation_mode(
domain.boundaries))
fluiddomain = FluidDomain(domain,
active=active_mask,
accessible=accessible_mask)
# --- Boundary Conditions, Pressure Solve ---
velocity = fluiddomain.with_hard_boundary_conditions(velocity)
for obstacle in obstacles:
if not obstacle.is_stationary:
obs_mask = mask(obstacle.geometry, antialias=True)
angular_velocity = AngularVelocity(
location=obstacle.geometry.center,
strength=obstacle.angular_velocity,
falloff=None)
velocity = ((1 - obs_mask) * velocity + obs_mask *
(angular_velocity + obstacle.velocity)).at(velocity)
divergence_field = velocity.divergence(physical_units=False)
if not struct.any(Material.open(
domain.boundaries)): # has no open boundary
divergence_field = divergence_field - math.mean(
divergence_field.data,
axis=tuple(range(1, 1 + divergence_field.rank)),
keepdims=True) # Subtract mean divergence
pressure, iterations = solve_pressure(divergence_field,
fluiddomain,
pressure_solver=pressure_solver)
pressure *= velocity.dx[0]
gradp = StaggeredGrid.gradient(pressure)
velocity -= fluiddomain.with_hard_boundary_conditions(gradp)
return velocity if not return_info else (velocity, {
'pressure': pressure,
'iterations': iterations,
'divergence': divergence_field
})
def _plot(axis, data, show_color_bar, vmin, vmax, **plt_args):
if isinstance(data, Grid) and data.spatial_rank == 1:
x = data.points.staggered_direction[0].vector[0].numpy()
requires_legend = False
for c in channel(data).meshgrid(names=True):
label = ", ".join([i for dim, i in c.items() if isinstance(i, str)])
values = data.values[c].numpy()
if values.dtype in (np.complex64, np.complex128):
axis.plot(x, values.real, label=f"real({label})" if label else "real")
axis.plot(x, values.imag, label=f"imag({label})" if label else "real")
requires_legend = True
else:
axis.plot(x, values, label=label)
requires_legend = requires_legend or label
if requires_legend:
axis.legend()
elif isinstance(data, Grid) and channel(data).volume == 1 and data.spatial_rank == 2:
dims = spatial(data)
if data.bounds.upper.vector.item_names is not None:
left, bottom = data.bounds.lower.vector[dims]
right, top = data.bounds.upper.vector[dims]
else:
dim_indices = data.resolution.indices(dims)
left, bottom = data.bounds.lower.vector[dim_indices]
right, top = data.bounds.upper.vector[dim_indices]
extent = (float(left), float(right), float(bottom), float(top))
im = axis.imshow(data.values.numpy(dims.reversed), origin='lower', extent=extent, vmin=vmin, vmax=vmax, **plt_args)
if show_color_bar:
axis.figure.colorbar(im, ax=axis) # adds a new Axis to the figure
axis.set_xlabel(dims.names[0])
axis.set_ylabel(dims.names[1])
elif isinstance(data, Grid) and data.spatial_rank == 2: # vector field
if isinstance(data, StaggeredGrid):
data = data.at_centers()
x, y = [d.numpy('x,y') for d in data.points.vector.unstack_spatial('x,y')]
u, v = [d.numpy('x,y') for d in data.values.vector.unstack_spatial('x,y')]
color = axis.xaxis.label.get_color()
axis.quiver(x, y, u, v, color=color, units='xy', scale=1)
axis.set_aspect('equal', adjustable='box')
elif isinstance(data, Grid) and channel(data).volume > 1 and data.spatial_rank == 3:
x, y, z = [d.numpy('x,y,z') for d in data.points.vector.unstack_spatial('x,y,z')]
u, v, w = [d.numpy('x,y,z') for d in data.values.vector.unstack_spatial('x,y,z')]
axis.quiver(x, y, z, u, v, w)
axis.set_xlabel('x')
axis.set_ylabel('y')
axis.set_zlabel('z')
elif isinstance(data, Grid) and channel(data).volume == 1 and data.spatial_rank == 3:
x, y, z = [d.numpy('x,y,z') for d in data.points.vector.unstack_spatial('x,y,z')]
values = data.values.numpy('x,y,z')
cmap = plt.get_cmap('viridis')
norm = matplotlib.colors.Normalize(vmin=np.min(values), vmax=np.max(values))
colors = cmap(norm(values))
axis.voxels(values, facecolors=colors, edgecolor='k')
elif isinstance(data, PointCloud) and data.spatial_rank == 2 and 'vector' in channel(data):
axis.set_aspect('equal', adjustable='box')
vector = data.points.shape['vector']
x, y = math.reshaped_native(data.points, [vector, data.shape.without('vector')], to_numpy=True, force_expand=True)
u, v = math.reshaped_native(data.values, [vector, data.shape.without('vector')], to_numpy=True, force_expand=True)
lower_x, lower_y = [float(d) for d in data.bounds.lower.vector]
upper_x, upper_y = [float(d) for d in data.bounds.upper.vector]
axis.set_xlim((lower_x, upper_x))
axis.set_ylim((lower_y, upper_y))
if data.color.shape:
color = data.color.numpy(data.shape.non_channel).reshape(-1)
else:
color = data.color.native()
axis.quiver(x, y, u, v, color=color, units='xy', scale=1)
if data.points.vector.item_names:
axis.set_xlabel(data.points.vector.item_names[0])
axis.set_ylabel(data.points.vector.item_names[1])
elif isinstance(data, PointCloud) and data.spatial_rank == 2:
axis.set_aspect('equal', adjustable='box')
lower_x, lower_y = [float(d) for d in data.bounds.lower.vector]
upper_x, upper_y = [float(d) for d in data.bounds.upper.vector]
axis.set_xlim((lower_x, upper_x))
axis.set_ylim((lower_y, upper_y))
if data.points.shape.non_channel.rank > 1: # multiple instance / spatial dimensions
data_list = field.unstack(data, data.points.shape.non_channel[0].name)
for d in data_list:
_plot_points(axis, d, **plt_args)
else:
_plot_points(axis, data, **plt_args)
elif isinstance(data, PointCloud) and data.spatial_rank == 3:
if data.points.shape.non_channel.rank > 1:
data_list = field.unstack(data, data.points.shape.non_channel[0].name)
for d in data_list:
_plot(axis, d, show_color_bar, vmin, vmax, **plt_args)
else:
x, y, z = [d.numpy() for d in data.points.vector.unstack_spatial('x,y,z')]
color = [d.native() for d in data.color.points.unstack(len(x))]
M = axis.transData.get_matrix()
x_scale, y_scale, z_scale = M[0, 0], M[1, 1], M[2, 2]
if isinstance(data.elements, Sphere):
symbol = 'o'
size = data.elements.bounding_radius().numpy() * 0.4
elif isinstance(data.elements, BaseBox):
symbol = 's'
size = math.mean(data.elements.bounding_half_extent(), 'vector').numpy() * 0.35
elif isinstance(data.elements, Point):
symbol = 'x'
size = 6 / (0.5 * (x_scale+y_scale+z_scale)/3)
else:
symbol = 'X'
size = data.elements.bounding_radius().numpy()
axis.scatter(x, y, z, marker=symbol, color=color, s=(size * 0.5 * (x_scale+y_scale+z_scale)/3) ** 2)
lower_x, lower_y, lower_z = [float(d) for d in data.bounds.lower.vector.unstack_spatial('x,y,z')]
upper_x, upper_y, upper_z = [float(d) for d in data.bounds.upper.vector.unstack_spatial('x,y,z')]
axis.set_xlim((lower_x, upper_x))
axis.set_ylim((lower_y, upper_y))
axis.set_zlim((lower_z, upper_z))
else:
raise NotImplementedError(f"No figure recipe for {data}")
def test_mean_collapsed(self):
ones = math.ones(x=40000, y=30000)
data = math.spatial_stack([ones, ones * 2], 'vector')
self.assertEqual(1.5, math.mean(data))
def _plot(data: SampledField,
fig: graph_objects.Figure,
size: tuple,
colormap: str or None,
show_color_bar: bool,
row: int = None, col: int = None,
):
subplot = fig.get_subplot(row, col)
vector = data.points.shape['vector']
if data.spatial_rank == 1 and isinstance(data, Grid):
x = data.points.vector[0].numpy().flatten()
channels = data.values.shape.channel
if channels.rank == 1 and channels.get_item_names(0) is not None:
for i, name in enumerate(channels.get_item_names(0)):
y = math.reshaped_native(real_values(data[{channels.name: i}]), [data.shape.spatial], to_numpy=True)
fig.add_trace(graph_objects.Scatter(x=x, y=y, mode='lines+markers', name=name), row=row, col=col)
fig.update_layout(showlegend=True)
else:
for ch_idx in channels.meshgrid():
y = math.reshaped_native(real_values(data[ch_idx]), [data.shape.spatial], to_numpy=True)
fig.add_trace(graph_objects.Scatter(x=x, y=y, mode='lines+markers', name='Multi-channel'), row=row, col=col)
fig.update_layout(showlegend=False)
elif data.spatial_rank == 2 and isinstance(data, Grid) and 'vector' not in data.shape: # heatmap
dims = spatial(data)
values = real_values(data).numpy(dims.reversed)
x = data.points.vector[dims[0].name].dimension(dims[1].name)[0].numpy()
y = data.points.vector[dims[1].name].dimension(dims[0].name)[0].numpy()
min_val, max_val = numpy.nanmin(values), numpy.nanmax(values)
min_val, max_val = min_val if numpy.isfinite(min_val) else 0, max_val if numpy.isfinite(max_val) else 0
color_scale = get_div_map(min_val, max_val, equal_scale=True, colormap=colormap)
# color_bar = graph_objects.heatmap.ColorBar(x=1.15) , colorbar=color_bar
fig.add_heatmap(row=row, col=col, x=x, y=y, z=values, zauto=False, zmin=min_val, zmax=max_val, colorscale=color_scale, showscale=show_color_bar)
subplot.xaxis.update(scaleanchor=f'y{subplot.yaxis.plotly_name[5:]}', scaleratio=1, constrain='domain', title=dims.names[0])
subplot.yaxis.update(constrain='domain', title=dims.names[1])
elif data.spatial_rank == 2 and isinstance(data, Grid): # vector field
if isinstance(data, StaggeredGrid):
data = data.at_centers()
x, y = math.reshaped_native(data.points.vector[spatial(data)], [vector, data.shape.without(vector)], to_numpy=True, force_expand=True)
extra_channels = data.shape.channel.without('vector')
data_x, data_y = math.reshaped_native(data.values, [vector, extra_channels, spatial(data)], to_numpy=True, force_expand=True)
lower_x, lower_y = [float(l) for l in data.bounds.lower.vector.unstack_spatial('x,y')]
upper_x, upper_y = [float(u) for u in data.bounds.upper.vector.unstack_spatial('x,y')]
x_range = [lower_x, upper_x]
y_range = [lower_y, upper_y]
for ch in range(data_x.shape[0]):
# quiver = figure_factory.create_quiver(x, y, data_x[ch], data_y[ch], scale=1.0) # 7 points per arrow
# fig.add_trace(quiver, row=row, col=col)
data_y_flat = data_y[ch].flatten()
data_x_flat = data_x[ch].flatten()
# lines_y = numpy.stack([y, y + data_y_flat, [None] * len(x)], -1).flatten() # 3 points per arrow
# lines_x = numpy.stack([x, x + data_x_flat, [None] * len(x)], -1).flatten()
lines_y = numpy.stack([y - data_y_flat / 2, y + data_y_flat / 2, [None] * len(x)], -1).flatten() # 3 points per arrow
lines_x = numpy.stack([x - data_x_flat / 2, x + data_x_flat / 2, [None] * len(x)], -1).flatten()
name = extra_channels.get_item_names(0)[ch] if extra_channels.rank == 1 and extra_channels.get_item_names(0) is not None else None
fig.add_scatter(x=lines_x, y=lines_y, mode='lines', row=row, col=col, name=name)
if data_x.shape[0] == 1:
fig.update_layout(showlegend=False)
subplot.xaxis.update(range=x_range)
subplot.yaxis.update(range=y_range)
subplot.xaxis.update(scaleanchor=f'y{subplot.yaxis.plotly_name[5:]}', scaleratio=1, constrain='domain')
subplot.yaxis.update(constrain='domain')
elif data.spatial_rank == 3 and isinstance(data, Grid) and data.shape.channel.volume == 1: # 3D heatmap
values = real_values(data).numpy('z,y,x')
x = data.points.vector['x'].numpy('z,y,x')
y = data.points.vector['y'].numpy('z,y,x')
z = data.points.vector['z'].numpy('z,y,x')
min_val, max_val = numpy.nanmin(values), numpy.nanmax(values)
min_val, max_val = min_val if numpy.isfinite(min_val) else 0, max_val if numpy.isfinite(max_val) else 0
color_scale = get_div_map(min_val, max_val, equal_scale=True, colormap=colormap)
fig.add_volume(x=x.flatten(), y=y.flatten(), z=z.flatten(), value=values.flatten(),
showscale=show_color_bar, colorscale=color_scale, cmin=min_val, cmax=max_val, cauto=False,
isomin=0.1, isomax=0.8,
opacity=0.1, # needs to be small to see through all surfaces
surface_count=17, # needs to be a large number for good volume rendering
row=row, col=col)
fig.update_layout(uirevision=True)
elif data.spatial_rank == 3 and isinstance(data, Grid): # 3D vector field
if isinstance(data, StaggeredGrid):
data = data.at_centers()
u = real_values(data).vector['x'].numpy('z,y,x')
v = real_values(data).vector['y'].numpy('z,y,x')
w = real_values(data).vector['z'].numpy('z,y,x')
x = data.points.vector['x'].numpy('z,y,x')
y = data.points.vector['y'].numpy('z,y,x')
z = data.points.vector['z'].numpy('z,y,x')
fig.add_cone(x=x.flatten(), y=y.flatten(), z=z.flatten(), u=u.flatten(), v=v.flatten(), w=w.flatten(),
colorscale='Blues',
sizemode="absolute", sizeref=1,
row=row, col=col)
elif isinstance(data, PointCloud) and data.spatial_rank == 2 and 'vector' in channel(data):
x, y = math.reshaped_native(data.points, [vector, data.shape.without('vector')], to_numpy=True, force_expand=True)
u, v = math.reshaped_native(data.values, [vector, data.shape.without('vector')], to_numpy=True, force_expand=True)
lower_x, lower_y = [float(d) for d in data.bounds.lower.vector]
upper_x, upper_y = [float(d) for d in data.bounds.upper.vector]
subplot.xaxis.update(range=[lower_x, upper_x])
subplot.yaxis.update(range=[lower_y, upper_y])
quiver = figure_factory.create_quiver(x, y, u, v, scale=1.0).data[0] # 7 points per arrow
if data.color.shape:
# color = data.color.numpy(data.shape.non_channel).reshape(-1)
warnings.warn("Multi-colored vector plots not yet supported")
else:
color = data.color.native()
quiver.line.update(color=color)
fig.add_trace(quiver, row=row, col=col)
if data.points.vector.item_names:
subplot.xaxis.update(title=data.points.vector.item_names[0])
subplot.yaxis.update(title=data.points.vector.item_names[1])
subplot.xaxis.update(scaleanchor=f'y{subplot.yaxis.plotly_name[5:]}', scaleratio=1, constrain='domain')
subplot.yaxis.update(constrain='domain')
elif isinstance(data, PointCloud) and data.spatial_rank == 2:
lower_x, lower_y = [float(d) for d in data.bounds.lower.vector]
upper_x, upper_y = [float(d) for d in data.bounds.upper.vector]
if data.points.shape.non_channel.rank > 1:
data_list = field.unstack(data, data.points.shape.non_channel[0].name)
for d in data_list:
_plot(d, fig, size, colormap, show_color_bar, row, col)
else:
x, y = [d.numpy() for d in data.points.vector.unstack_spatial('x,y')]
color = data.color.native()
subplot_height = (subplot.yaxis.domain[1] - subplot.yaxis.domain[0]) * size[1]
if isinstance(data.elements, Sphere):
symbol = 'circle'
marker_size = data.elements.bounding_radius().numpy() * 1.9
elif isinstance(data.elements, BaseBox):
symbol = 'square'
marker_size = math.mean(data.elements.bounding_half_extent(), 'vector').numpy() * 1
elif isinstance(data.elements, Point):
symbol = 'x'
marker_size = 12 / (subplot_height / (upper_y - lower_y))
else:
symbol = 'asterisk'
marker_size = data.elements.bounding_radius().numpy()
marker_size *= subplot_height / (upper_y - lower_y)
marker = graph_objects.scatter.Marker(size=marker_size, color=color, sizemode='diameter', symbol=symbol)
fig.add_scatter(mode='markers', x=x, y=y, marker=marker, row=row, col=col)
subplot.xaxis.update(range=[lower_x, upper_x])
subplot.yaxis.update(range=[lower_y, upper_y])
fig.update_layout(showlegend=False)
subplot.xaxis.update(scaleanchor=f'y{subplot.yaxis.plotly_name[5:]}', scaleratio=1, constrain='domain')
subplot.yaxis.update(constrain='domain')
elif isinstance(data, PointCloud) and data.spatial_rank == 3:
lower_x, lower_y, lower_z = [float(d) for d in data.bounds.lower.vector.unstack_spatial('x,y,z')]
upper_x, upper_y, upper_z = [float(d) for d in data.bounds.upper.vector.unstack_spatial('x,y,z')]
if data.points.shape.non_channel.rank > 1:
data_list = field.unstack(data, data.points.shape.non_channel[0].name)
for d in data_list:
_plot(d, fig, size, colormap, show_color_bar, row, col)
else:
x, y, z = [d.numpy() for d in data.points.vector.unstack_spatial('x,y,z')]
color = data.color.native()
# if data.color.shape.instance_rank == 0:
# color = str(data.color)
# else:
# color = [str(d) for d in math.unstack(data.color, instance)]
domain_y = fig.layout[subplot.plotly_name].domain.y
if isinstance(data.elements, Sphere):
symbol = 'circle'
marker_size = data.elements.bounding_radius().numpy() * 2
elif isinstance(data.elements, BaseBox):
symbol = 'square'
marker_size = math.mean(data.elements.bounding_half_extent(), 'vector').numpy() * 1
elif isinstance(data.elements, Point):
symbol = 'x'
marker_size = 4 / (size[1] * (domain_y[1] - domain_y[0]) / (upper_y - lower_y) * 0.5)
else:
symbol = 'asterisk'
marker_size = data.elements.bounding_radius().numpy()
marker_size *= size[1] * (domain_y[1] - domain_y[0]) / (upper_y - lower_y) * 0.5
marker = graph_objects.scatter3d.Marker(size=marker_size, color=color, sizemode='diameter', symbol=symbol)
fig.add_scatter3d(mode='markers', x=x, y=y, z=z, marker=marker, row=row, col=col)
subplot.xaxis.update(range=[lower_x, upper_x])
subplot.yaxis.update(range=[lower_y, upper_y])
subplot.zaxis.update(range=[lower_z, upper_z])
fig.update_layout(showlegend=False)
else:
raise NotImplementedError(f"No figure recipe for {data}")
def mean(field: Grid):
return math.mean(field.values, field.shape.spatial)
def test_mean_collapsed(self):
ones = math.ones(spatial(x=40000, y=30000))
data = math.stack([ones, ones * 2], spatial('vector'))
math.assert_close(1.5, math.mean(data))
