def batch_stack(*fields, dim: str):
assert all(isinstance(f, SampledField) for f in fields)
assert all(isinstance(f, type(fields[0])) for f in fields)
if any(f.extrapolation != fields[0].extrapolation for f in fields):
raise NotImplementedError("Concatenating extrapolations not supported")
if isinstance(fields[0], Grid):
values = math.batch_stack([f.values for f in fields], dim)
return fields[0].with_(values=values)
elif isinstance(fields[0], PointCloud):
elements = geom.stack(*[f.elements for f in fields], dim=dim)
values = math.batch_stack([f.values for f in fields], dim=dim)
colors = math.batch_stack([f.color for f in fields], dim=dim)
return fields[0].with_(elements=elements, values=values, color=colors)
raise NotImplementedError(type(fields[0]))
def lies_inside(self, location: math.Tensor):
if self.stack_dim_name in location.shape:
location = location.unstack(self.stack_dim_name)
else:
location = [location] * len(self.geometries)
inside = [
g.lies_inside(loc) for g, loc in zip(self.geometries, location)
]
return math.batch_stack(inside, self.stack_dim_name)
def test_tensor_from_tensor(self):
ref = math.batch_stack([math.zeros(x=5), math.zeros(x=4)], 'stack')
for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
with backend:
tens = math.tensor(ref, convert=False)
self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens))
self.assertEqual(2, tens.shape.get_size('stack'))
self.assertEqual(('stack', 'x'), tens.shape.names)
tens = math.tensor(ref)
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(backend, math.choose_backend(tens.stack[0]))
self.assertEqual(backend, math.choose_backend(tens.stack[1]))
tens = math.tensor(ref, names=('n1', 'n2'))
self.assertEqual(backend, math.choose_backend(tens))
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor:
if not reduce_channels:
if geometry == self.elements:
return self.values
elif isinstance(geometry, GridCell):
return self._grid_scatter(geometry.bounds, geometry.resolution)
elif isinstance(geometry, GeometryStack):
sampled = [self.sample_at(g) for g in geometry.geometries]
return math.batch_stack(sampled, geometry.stack_dim_name)
else:
raise NotImplementedError()
else:
assert len(reduce_channels) == 1
components = self.unstack('vector') if 'vector' in self.shape else (self,) * geometry.shape.get_size(reduce_channels[0])
sampled = [c.sample_in(p) for c, p in zip(components, geometry.unstack(reduce_channels[0]))]
return math.channel_stack(sampled, 'vector')
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor:
if reduce_channels:
assert len(reduce_channels) == 1
geometries = geometry.unstack(reduce_channels[0])
components = self.vector.unstack(len(geometries))
sampled = [c.sample_in(g) for c, g in zip(components, geometries)]
return math.channel_stack(sampled, 'vector')
if isinstance(geometry, GeometryStack):
sampled = [self.sample_in(g) for g in geometry.geometries]
return math.batch_stack(sampled, geometry.stack_dim_name)
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
return self.sample_at(geometry.center, reduce_channels)
def batch_stack(*scenes: 'Scene', dim: str = 'batch') -> 'Scene':
return Scene(math.batch_stack([s._paths for s in scenes], dim))
def bounding_half_extent(self):
values = [g.bounding_half_extent() for g in self.geometries]
return math.batch_stack(values, self.stack_dim_name)
def bounding_radius(self):
radii = [g.bounding_radius() for g in self.geometries]
return math.batch_stack(radii, self.stack_dim_name)
def center(self):
centers = [g.center for g in self.geometries]
return math.batch_stack(centers, self.stack_dim_name)