import sys
import os
import numpy as np
import bpy
from enum import Enum, Flag, auto
from typing import Union, List, Optional
from goo.handler import Handler, StopHandler
from goo.cell import Cell, CellType
from goo.molecule import DiffusionSystem
Render = Enum("Render", ["PNG", "TIFF", "MP4"])
[docs]
class Simulator:
"""A simulator for cell-based simulations in Blender.
Args:
cells (List[Cell]): List of cells.
time (List[int]): Start and end frames.
physics_dt (int): Time step for physics simulation.
molecular_dt (int): Time step for molecular simulation.
"""
# TODO: determine diffsystem or diffsystems
def __init__(
self,
celltypes: List[Union[CellType, Cell]] = [],
diffsystems: DiffusionSystem = [],
time: int = 250,
physics_dt: int = 1,
molecular_dt: int = 1,
):
self.celltypes = celltypes
# takes first possible diffsystem
self.diffsystem = diffsystems[0] if diffsystems else None
self.physics_dt = physics_dt
self.molecular_dt = molecular_dt
self.addons = ["add_mesh_extra_objects"]
self.render_format: Render = Render.PNG
self.time = time
# Set up simulation parameters for diffusion system
if self.diffsystem is not None:
self.diffsystem.time_step = molecular_dt
self.diffsystem.total_time = physics_dt
[docs]
def set_seed(self, seed):
"""Set the random seed for the simulation."""
np.random.seed(seed)
bpy.context.scene["seed"] = seed
[docs]
def setup_world(self, seed=1):
"""Set up the Blender scene for the simulation."""
# Enable addons
for addon in self.addons:
self.enable_addon(addon)
# Set random seed
self.set_seed(seed)
# Set up simulation time interval
bpy.context.scene.frame_start = 1
bpy.context.scene.frame_end = self.time
# Set units to the metric system
bpy.context.scene.unit_settings.system = "METRIC"
bpy.context.scene.unit_settings.scale_length = 1e-6
bpy.context.scene.unit_settings.system_rotation = "DEGREES"
bpy.context.scene.unit_settings.length_unit = "MICROMETERS"
bpy.context.scene.unit_settings.mass_unit = "MILLIGRAMS"
bpy.context.scene.unit_settings.time_unit = "SECONDS"
bpy.context.scene.unit_settings.temperature_unit = "CELSIUS"
# Turn off gravity
self.toggle_gravity(False)
# Set up rendering environment
node_tree = bpy.context.scene.world.node_tree
tree_nodes = node_tree.nodes
tree_nodes.clear()
# Add background node
node_background = tree_nodes.new(type="ShaderNodeBackground")
node_environment = tree_nodes.new("ShaderNodeTexEnvironment")
scripts_paths = bpy.utils.script_paths()
try:
node_environment.image = bpy.data.images.load(
scripts_paths[-1] + "/modules/goo/missile_launch_facility_01_4k.hdr"
)
except Exception:
print(sys.exc_info())
print(
"""WARNING FROM GOO: To enable proper rendering you must have
/modules/goo/missile_launch_facility_01_4k.hdr in the right location"""
)
node_environment.location = -300, 0
# Add output node
node_output = tree_nodes.new(type="ShaderNodeOutputWorld")
node_output.location = 200, 0
# Link all nodes
links = node_tree.links
links.new(node_environment.outputs["Color"], node_background.inputs["Color"])
links.new(node_background.outputs["Background"], node_output.inputs["Surface"])
# # set film to transparent to hide background
bpy.context.scene.render.film_transparent = True
# Allow cloth physics pass 250 frames
self.extend_scene()
[docs]
def enable_addon(self, addon):
"""Enable an addon in Blender."""
if addon not in bpy.context.preferences.addons:
bpy.ops.preferences.addon_enable(module=addon)
print(f"Addon '{addon}' has been enabled.")
else:
print(f"Addon '{addon}' is already enabled.")
[docs]
def toggle_gravity(self, on):
"""Toggle gravity in the scene."""
bpy.context.scene.use_gravity = on
[docs]
def get_cells_func(self, celltypes=None):
"""Get a function that returns all cells in the simulation."""
celltypes = celltypes if celltypes else self.celltypes
def get_cells():
return [cell for celltype in celltypes for cell in celltype.cells]
return get_cells
[docs]
def get_diffsystem_func(self, diffsystem=None):
"""Get a function that returns the diffusion system."""
diffsystem = diffsystem if diffsystem is not None else self.diffsystem
def get_diffsystem():
return diffsystem
return get_diffsystem
[docs]
def get_cells(self, celltypes=None):
"""Get all cells in the simulation."""
celltypes = celltypes if celltypes else self.celltypes
return [cell for celltype in celltypes for cell in celltype.cells]
[docs]
def extend_scene(self):
"""Extend the scene to allow cloth physics to pass the default 250 frames."""
cells = self.get_cells()
for cell in cells:
if cell.cloth_mod and cell.cloth_mod.point_cache.frame_end < self.time:
cell.cloth_mod.point_cache.frame_end = self.time
[docs]
def add_handler(
self,
handler: Handler,
celltypes: list[CellType] = None,
diffsystem: DiffusionSystem = None,
):
"""Add a handler to the simulation."""
handler.setup(
self.get_cells_func(celltypes),
self.get_diffsystem_func(diffsystem),
self.physics_dt,
)
bpy.app.handlers.frame_change_post.append(handler.run)
[docs]
def add_handlers(self, handlers: list[Handler]):
"""Add multiple handlers to the simulation."""
# always include a stop handler
stop_handler = StopHandler()
bpy.app.handlers.frame_change_pre.append(stop_handler.run)
for handler in handlers:
self.add_handler(handler)
[docs]
def render(
self,
frames: Optional[Union[List[int], range]] = None,
path: str = None,
camera=False,
format: Render = Render.PNG,
):
"""
Render the simulation in the background without
updating the 3D Viewport in real time.
If a camera is specified, the frames will be rendered with it,
otherwise the frames will be rendered in the 3D Viewport.
It will updated the scene at the end of the simulation.
Args:
start (int): Start frame.
end (int): End frame.
path (str): Path to save the frames.
camera (bool): Render with the camera.
format (Render): Render format: PNG (default), TIFF, MP4.
"""
if not path:
path = os.path.dirname(bpy.context.scene.render.filepath)
else:
print("Save path not provided. Falling back on default path.")
print("----- RENDERING... -----")
render_format = format if format else self.render_format
match render_format:
case Render.PNG:
bpy.context.scene.render.image_settings.file_format = "PNG"
case Render.TIFF:
bpy.context.scene.render.image_settings.file_format = "TIFF"
case Render.MP4:
bpy.context.scene.render.image_settings.file_format = "FFMPEG"
bpy.context.scene.render.ffmpeg.format = "MPEG4"
match frames:
case None:
start = 1
end = self.time
frame_list = list(range(start, end + 1))
case range():
frame_list = list(frames)
case list():
frame_list = frames
for i in range(1, max(frame_list) + 1):
print(i, end=" ")
bpy.context.scene.frame_set(i)
bpy.context.scene.render.filepath = os.path.join(path, f"{i:04d}")
if i in frame_list:
if camera:
bpy.ops.render.render(write_still=True)
else:
bpy.ops.render.opengl(write_still=True)
bpy.context.scene.render.filepath = path
print("\n----- RENDERING COMPLETED! -----")
[docs]
def render_animation(self, path=None, end=bpy.context.scene.frame_end, camera=False):
"""Render the simulation as an animation."""
if not path:
print("Save path not provided. Falling back on default path.")
path = os.path.dirname(bpy.context.scene.render.filepath)
bpy.context.scene.render.filepath = os.path.join(path, "")
bpy.context.scene.render.image_settings.file_format = "FFMPEG"
bpy.context.scene.render.ffmpeg.format = "MPEG4"
print("----- RENDERING... -----")
print("Rendering to", bpy.context.scene.render.filepath)
bpy.context.scene.frame_start = 1
bpy.context.scene.frame_set(1)
bpy.context.scene.frame_end = end
if camera:
bpy.ops.render.render(animation=True, write_still=True)
else:
bpy.ops.render.opengl(animation=True, write_still=True)
print("\n----- RENDERING COMPLETED! -----")
[docs]
def run(self, end=bpy.context.scene.frame_end):
"""
Run the simulation in the background without
updating the 3D Viewport in real time.
Args:
end (int): End frame. Defaults to the last frame of the scene.
"""
print("----- SIMULATION START -----")
for i in range(1, end + 1):
print(i, end=" ")
bpy.context.scene.frame_set(i)
print("\n----- SIMULATION END -----")