tuned pipeline to 60 fps
This commit is contained in:
@@ -57,6 +57,11 @@ class MaskedDomain:
|
||||
"num_active",
|
||||
"pinned",
|
||||
"active_coords",
|
||||
"active_linear_indices",
|
||||
"_rhs_buffer",
|
||||
"_rhs_pair_buffer",
|
||||
"_scatter_buffer",
|
||||
"_scatter_pair_buffer",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
@@ -69,6 +74,7 @@ class MaskedDomain:
|
||||
num_active,
|
||||
pinned,
|
||||
active_coords,
|
||||
active_linear_indices,
|
||||
):
|
||||
self.shape = shape
|
||||
self.mask = mask
|
||||
@@ -78,6 +84,11 @@ class MaskedDomain:
|
||||
self.num_active = num_active
|
||||
self.pinned = pinned
|
||||
self.active_coords = active_coords
|
||||
self.active_linear_indices = active_linear_indices
|
||||
self._rhs_buffer = None
|
||||
self._rhs_pair_buffer = None
|
||||
self._scatter_buffer = None
|
||||
self._scatter_pair_buffer = None
|
||||
|
||||
@classmethod
|
||||
def build(cls, mask, spacings):
|
||||
@@ -101,10 +112,12 @@ class MaskedDomain:
|
||||
num_active=0,
|
||||
pinned=0,
|
||||
active_coords=active_coords,
|
||||
active_linear_indices=numpy.flatnonzero(active),
|
||||
)
|
||||
|
||||
index_map = numpy.full(active.shape, -1, dtype=numpy.int32)
|
||||
index_map[active] = numpy.arange(num_active, dtype=numpy.int32)
|
||||
active_linear_indices = numpy.flatnonzero(active)
|
||||
|
||||
dy, dx = float(spacings[0]), float(spacings[1])
|
||||
coeff_y = 1.0 / (dy * dy)
|
||||
@@ -153,6 +166,7 @@ class MaskedDomain:
|
||||
num_active=num_active,
|
||||
pinned=pinned,
|
||||
active_coords=active_coords,
|
||||
active_linear_indices=active_linear_indices,
|
||||
)
|
||||
|
||||
def matches(self, mask):
|
||||
@@ -166,13 +180,34 @@ class MaskedDomain:
|
||||
f: (H, W) scalar source field
|
||||
returns: (N,) RHS on active pixels; empty if no active pixels
|
||||
"""
|
||||
rhs = numpy.asarray(f, dtype=numpy.float64)[self.mask]
|
||||
if rhs.size == 0:
|
||||
return rhs
|
||||
if self.num_active == 0:
|
||||
return numpy.zeros(0, dtype=numpy.float64)
|
||||
if self._rhs_buffer is None or self._rhs_buffer.shape != (self.num_active,):
|
||||
self._rhs_buffer = numpy.empty(self.num_active, dtype=numpy.float64)
|
||||
numpy.take(
|
||||
numpy.asarray(f, dtype=numpy.float64).ravel(),
|
||||
self.active_linear_indices,
|
||||
out=self._rhs_buffer,
|
||||
)
|
||||
rhs = self._rhs_buffer
|
||||
rhs -= rhs.mean()
|
||||
rhs[self.pinned] = 0.0
|
||||
return rhs
|
||||
|
||||
def prepare_rhs_pair(self, f0, f1):
|
||||
"""Extract two gauge-fixed RHS vectors into one reusable (N, 2) buffer."""
|
||||
if self.num_active == 0:
|
||||
return numpy.zeros((0, 2), dtype=numpy.float64)
|
||||
if self._rhs_pair_buffer is None or self._rhs_pair_buffer.shape != (self.num_active, 2):
|
||||
self._rhs_pair_buffer = numpy.empty((self.num_active, 2), dtype=numpy.float64)
|
||||
flat0 = numpy.asarray(f0, dtype=numpy.float64).ravel()
|
||||
flat1 = numpy.asarray(f1, dtype=numpy.float64).ravel()
|
||||
numpy.take(flat0, self.active_linear_indices, out=self._rhs_pair_buffer[:, 0])
|
||||
numpy.take(flat1, self.active_linear_indices, out=self._rhs_pair_buffer[:, 1])
|
||||
self._rhs_pair_buffer -= self._rhs_pair_buffer.mean(axis=0, keepdims=True)
|
||||
self._rhs_pair_buffer[self.pinned, :] = 0.0
|
||||
return self._rhs_pair_buffer
|
||||
|
||||
def scatter(self, sol_active, template):
|
||||
"""Map reduced solution back to full grid.
|
||||
|
||||
@@ -180,10 +215,25 @@ class MaskedDomain:
|
||||
template: (H, W) array used for output shape/dtype
|
||||
returns: (H, W) potential; zero outside mask
|
||||
"""
|
||||
out = numpy.zeros_like(template, dtype=numpy.float32)
|
||||
if self._scatter_buffer is None or self._scatter_buffer.shape != template.shape:
|
||||
self._scatter_buffer = numpy.zeros(template.shape, dtype=numpy.float32)
|
||||
else:
|
||||
self._scatter_buffer.fill(0)
|
||||
if sol_active.size:
|
||||
out[self.mask] = numpy.asarray(sol_active, dtype=numpy.float32)
|
||||
return out
|
||||
self._scatter_buffer.ravel()[self.active_linear_indices] = numpy.asarray(sol_active, dtype=numpy.float32)
|
||||
return self._scatter_buffer.copy()
|
||||
|
||||
def scatter_pair(self, sol_active_pair, template):
|
||||
"""Map two reduced solutions back to a pair of full-grid outputs."""
|
||||
if self._scatter_pair_buffer is None or self._scatter_pair_buffer.shape != (2,) + template.shape:
|
||||
self._scatter_pair_buffer = numpy.zeros((2,) + template.shape, dtype=numpy.float32)
|
||||
else:
|
||||
self._scatter_pair_buffer.fill(0)
|
||||
if sol_active_pair.size:
|
||||
values = numpy.asarray(sol_active_pair, dtype=numpy.float32)
|
||||
self._scatter_pair_buffer[0].ravel()[self.active_linear_indices] = values[:, 0]
|
||||
self._scatter_pair_buffer[1].ravel()[self.active_linear_indices] = values[:, 1]
|
||||
return self._scatter_pair_buffer[0].copy(), self._scatter_pair_buffer[1].copy()
|
||||
|
||||
def residual(self, field, rhs_active):
|
||||
"""Compute masked Laplacian residual L u - rhs on active pixels.
|
||||
|
||||
Reference in New Issue
Block a user