Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Static mesh support #45

Closed
wants to merge 6 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions znvis/particle/particle.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,10 @@ class Particle:
Director tensor of the shape (n_confs, n_particles, n_dims)
mesh_list : list
A list of mesh objects, one for each time step.
static : bool (default=False)
If true, only render the mesh once at initialization. Be careful
as this changes the shape of the required position and director
to (n_particles, n_dims)
smoothing : bool (default=False)
If true, apply smoothing to each mesh object as it is rendered.
This will slow down the initial construction of the mesh objects
Expand All @@ -64,7 +68,7 @@ class Particle:
force: np.ndarray = None
director: np.ndarray = None
mesh_list: typing.List[Mesh] = None

static: bool = False
smoothing: bool = False

def _create_mesh(self, position, director):
Expand Down Expand Up @@ -105,13 +109,21 @@ def construct_mesh_list(self):
"""
self.mesh_list = []
try:
# n_particles = int(self.position.shape[1])
n_time_steps = int(len(self.position))
if not self.static:
n_particles = int(self.position.shape[1])
n_time_steps = int(self.position.shape[0])
else:
n_particles = int(self.position.shape[0])
n_time_steps = 1
self.position = self.position[np.newaxis, :, :]
if self.director is not None:
self.director = self.director[np.newaxis, :, :]

except ValueError:
raise ValueError("There is no data for these particles.")

for i in track(range(n_time_steps), description=f"Building {self.name} Mesh"):
for j in range(np.shape(self.position[i])[0]):
for j in range(n_particles):
if j == 0:
if self.director is not None:
mesh = self._create_mesh(
Expand Down
31 changes: 24 additions & 7 deletions znvis/particle/vector_field.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,10 @@ class VectorField:
Direction tensor of the shape (n_steps, n_vectors, n_dims)
mesh_list : list
A list of mesh objects, one for each time step.
static : bool (default=False)
If true, only render the mesh once at initialization. Be careful
as this changes the shape of the required position and direction
to (n_particles, n_dims)
smoothing : bool (default=False)
If true, apply smoothing to each mesh object as it is rendered.
This will slow down the initial construction of the mesh objects
Expand All @@ -58,7 +62,7 @@ class VectorField:
position: np.ndarray = None
direction: np.ndarray = None
mesh_list: typing.List[Arrow] = None

static: bool = False
smoothing: bool = False

def _create_mesh(self, position: np.ndarray, direction: np.ndarray):
Expand Down Expand Up @@ -97,15 +101,28 @@ def construct_mesh_list(self):
"""
self.mesh_list = []
try:
n_particles = int(self.position.shape[1])
n_time_steps = int(self.position.shape[0])
if not self.static:
n_particles = int(self.position.shape[1])
n_time_steps = int(self.position.shape[0])
else:
n_particles = int(self.position.shape[0])
n_time_steps = 1
self.position = self.position[np.newaxis, :, :]
self.direction = self.direction[np.newaxis, :, :]

except ValueError:
raise ValueError("There is no data for this vector field.")

new_mesh = True

for i in track(range(n_time_steps), description=f"Building {self.name} Mesh"):
for j in range(n_particles):
if j == 0:
mesh = self._create_mesh(self.position[i][j], self.direction[i][j])
else:
mesh += self._create_mesh(self.position[i][j], self.direction[i][j])
if np.max(self.direction[i][j]) > 0: # ignore vectors with length zero
if new_mesh is False:
mesh += self._create_mesh(self.position[i][j], self.direction[i][j])
else:
mesh = self._create_mesh(self.position[i][j], self.direction[i][j])
new_mesh = False
new_mesh = True

self.mesh_list.append(mesh)
104 changes: 77 additions & 27 deletions znvis/visualizer/visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,11 @@ def __init__(
self.bounding_box = bounding_box() if bounding_box else None

if number_of_steps is None:
number_of_steps = len(particles[0].position)
self.number_of_steps = number_of_steps
len_list = []
for particle in particles:
if not particle.static:
len_list.append(len(particle.position))
self.number_of_steps = min(len_list)

self.output_folder = pathlib.Path(output_folder).resolve()
self.frame_folder = self.output_folder / "video_frames"
Expand Down Expand Up @@ -271,18 +274,35 @@ def _take_screenshot(self, vis):
old_state = self.interrupt # get old state
self.interrupt = 0 # stop live feed if running.
mesh_dict = {}
mesh_center = []

if self.vector_field is not None:
for item in self.vector_field:
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

for item in self.particles:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
mesh_center.append(
item.mesh_list[self.counter]
.get_axis_aligned_bounding_box()
.get_center()
)
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

view_matrix = vis.scene.camera.get_view_matrix()

Expand Down Expand Up @@ -346,10 +366,11 @@ def _draw_particles(self, visualizer=None, initial: bool = False):
visualizer.add_geometry("Box", self.bounding_box)
else:
for i, item in enumerate(self.particles):
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)
if not item.static:
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)


def _draw_vector_field(self, visualizer=None, initial: bool = False):
Expand All @@ -376,10 +397,11 @@ def _draw_vector_field(self, visualizer=None, initial: bool = False):
)
else:
for i, item in enumerate(self.vector_field):
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)
if not item.static:
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)

def _continuous_trajectory(self, vis):
"""
Expand Down Expand Up @@ -415,12 +437,34 @@ def save_callable():
"""
mesh_dict = {}

if self.vector_field is not None:
for item in self.vector_field:
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

for item in self.particles:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

view_matrix = self.vis.scene.camera.get_view_matrix()
self.renderer.render_mesh_objects(
Expand Down Expand Up @@ -472,10 +516,16 @@ def save_callable():
Function to be called on thread to save image.
"""
for i, item in enumerate(self.particles):
if i == 0:
mesh = item.mesh_list[self.counter]
if item.static:
if i == 0:
mesh = item.mesh_list[0]
else:
mesh += item.mesh_list[0]
else:
mesh += item.mesh_list[self.counter]
if i == 0:
mesh = item.mesh_list[self.counter]
else:
mesh += item.mesh_list[self.counter]

o3d.io.write_triangle_mesh(
(self.obj_folder / f"export_mesh_{self.counter}.ply").as_posix(), mesh
Expand Down
Loading