feat: Proper visualization pathways

main
Sofus Albert Høgsbro Rose 2024-04-01 19:28:24 +02:00
parent 221d5378e4
commit e080d16893
Signed by: so-rose
GPG Key ID: AD901CB0F3701434
9 changed files with 520 additions and 125 deletions

View File

@ -89,27 +89,6 @@ def import_geonodes(
return bpy.data.node_groups[geonodes]
####################
# - GeoNodes Asset Shelf
####################
# class GeoNodesAssetShelf(bpy.types.AssetShelf):
# bl_space_type = 'NODE_EDITOR'
# bl_idname = 'blender_maxwell.asset_shelf__geonodes'
# bl_options = {'NO_ASSET_DRAG'}
#
# @classmethod
# def poll(cls, context):
# return (
# (space := context.get('space_data'))
# and (node_tree := space.get('node_tree'))
# and (node_tree.bl_idname == 'MaxwellSimTreeType')
# )
#
# @classmethod
# def asset_poll(cls, asset: bpy.types.AssetRepresentation):
# return asset.id_type == 'NODETREE'
####################
# - GeoNodes Asset Shelf Panel for MaxwellSimTree
####################

View File

@ -4,5 +4,10 @@ from ....utils.blender_type_enum import BlenderTypeEnum
class ManagedObjType(BlenderTypeEnum):
ManagedBLObject = enum.auto()
ManagedBLImage = enum.auto()
ManagedBLCollection = enum.auto()
ManagedBLEmpty = enum.auto()
ManagedBLMesh = enum.auto()
ManagedBLVolume = enum.auto()
ManagedBLModifier = enum.auto()

View File

@ -1,2 +1,8 @@
from .managed_bl_image import ManagedBLImage
from .managed_bl_object import ManagedBLObject
#from .managed_bl_collection import ManagedBLCollection
#from .managed_bl_object import ManagedBLObject
from .managed_bl_mesh import ManagedBLMesh
from .managed_bl_empty import ManagedBLEmpty
#from .managed_bl_volume import ManagedBLVolume
from .managed_bl_modifier import ManagedBLModifier

View File

@ -0,0 +1,39 @@
import bpy
from ....utils import logger
log = logger.get(__name__)
MANAGED_COLLECTION_NAME = 'BLMaxwell'
PREVIEW_COLLECTION_NAME = 'BLMaxwell Visible'
####################
# - Global Collection Handling
####################
def collection(collection_name: str, view_layer_exclude: bool) -> bpy.types.Collection:
# Init the "Managed Collection"
# Ensure Collection exists (and is in the Scene collection)
if collection_name not in bpy.data.collections:
collection = bpy.data.collections.new(collection_name)
bpy.context.scene.collection.children.link(collection)
else:
collection = bpy.data.collections[collection_name]
## Ensure synced View Layer exclusion
if (
layer_collection := bpy.context.view_layer.layer_collection.children[
collection_name
]
).exclude != view_layer_exclude:
layer_collection.exclude = view_layer_exclude
return collection
def managed_collection() -> bpy.types.Collection:
return collection(MANAGED_COLLECTION_NAME, view_layer_exclude=False)
def preview_collection() -> bpy.types.Collection:
return collection(PREVIEW_COLLECTION_NAME, view_layer_exclude=True)

View File

@ -0,0 +1,228 @@
import contextlib
import bmesh
import bpy
import numpy as np
import typing_extensions as typx
from ....utils import logger
from .. import contracts as ct
from .managed_bl_collection import managed_collection, preview_collection
log = logger.get(__name__)
####################
# - BLMesh
####################
class ManagedBLMesh(ct.schemas.ManagedObj):
managed_obj_type = ct.ManagedObjType.ManagedBLMesh
_bl_object_name: str | None = None
####################
# - BL Object Name
####################
@property
def name(self):
return self._bl_object_name
@name.setter
def name(self, value: str) -> None:
log.info(
'Changing BLMesh w/Name "%s" to Name "%s"', self._bl_object_name, value
)
if not bpy.data.objects.get(value):
log.info(
'Desired BLMesh Name "%s" Not Taken',
value,
)
if self._bl_object_name is None:
log.info(
'Set New BLMesh Name to "%s"',
value,
)
elif bl_object := bpy.data.objects.get(self._bl_object_name):
log.info(
'Changed BLMesh Name to "%s"',
value,
)
bl_object.name = value
else:
msg = f'ManagedBLMesh with name "{self._bl_object_name}" was deleted'
raise RuntimeError(msg)
# Set Internal Name
self._bl_object_name = value
else:
log.info(
'Desired BLMesh Name "%s" is Taken. Using Blender Rename',
value,
)
# Set Name Anyway, but Respect Blender's Renaming
## When a name already exists, Blender adds .### to prevent overlap.
## `set_name` is allowed to change the name; nodes account for this.
bl_object.name = value
self._bl_object_name = bl_object.name
log.info(
'Changed BLMesh Name to "%s"',
bl_object.name,
)
####################
# - Allocation
####################
def __init__(self, name: str):
self.name = name
####################
# - Deallocation
####################
def free(self):
if (bl_object := bpy.data.objects.get(self.name)) is None:
return
# Delete the Underlying Datablock
## This automatically deletes the object too
log.info('Removing "%s" BLMesh', bl_object.type)
bpy.data.meshes.remove(bl_object.data)
####################
# - Actions
####################
def show_preview(self) -> None:
"""Moves the managed Blender object to the preview collection.
If it's already included, do nothing.
"""
bl_object = self.bl_object()
if bl_object.name not in preview_collection().objects:
log.info('Moving "%s" to Preview Collection', bl_object.name)
preview_collection().objects.link(bl_object)
def hide_preview(self) -> None:
"""Removes the managed Blender object from the preview collection.
If it's already removed, do nothing.
"""
bl_object = self.bl_object()
if bl_object.name not in preview_collection().objects:
log.info('Removing "%s" from Preview Collection', bl_object.name)
preview_collection.objects.unlink(bl_object)
def bl_select(self) -> None:
"""Selects the managed Blender object, causing it to be ex. outlined in the 3D viewport."""
if (bl_object := bpy.data.objects.get(self.name)) is not None:
bpy.ops.object.select_all(action='DESELECT')
bl_object.select_set(True)
msg = 'Managed BLMesh does not exist'
raise ValueError(msg)
####################
# - BLMesh Management
####################
def bl_object(self):
"""Returns the managed blender object."""
# Create Object w/Appropriate Data Block
if not (bl_object := bpy.data.objects.get(self.name)):
log.info(
'Creating BLMesh Object "%s"',
bl_object.name,
)
bl_data = bpy.data.meshes.new(self.name)
bl_object = bpy.data.objects.new(self.name, bl_data)
log.debug(
'Linking "%s" to Base Collection',
bl_object.name,
)
managed_collection().objects.link(bl_object)
return bl_object
####################
# - Mesh Data Properties
####################
@property
def mesh_data(self) -> bpy.types.Mesh:
"""Directly loads the Blender mesh data.
Raises:
ValueError: If the object has no mesh data.
"""
if bl_object := bpy.data.objects.get(self.name):
return bl_object.data
msg = f'Requested mesh data from {self.name} of type {bl_object.type}'
raise ValueError(msg)
@contextlib.contextmanager
def mesh_as_bmesh(
self,
evaluate: bool = True,
triangulate: bool = False,
) -> bpy.types.Mesh:
if (bl_object := bpy.data.objects.get(self.name)) and bl_object.type == 'MESH':
bmesh_mesh = None
try:
bmesh_mesh = bmesh.new()
if evaluate:
bmesh_mesh.from_object(
bl_object,
bpy.context.evaluated_depsgraph_get(),
)
else:
bmesh_mesh.from_object(bl_object)
if triangulate:
bmesh.ops.triangulate(bmesh_mesh, faces=bmesh_mesh.faces)
yield bmesh_mesh
finally:
if bmesh_mesh:
bmesh_mesh.free()
else:
msg = f'Requested BMesh from "{self.name}" of type "{bl_object.type}"'
raise ValueError(msg)
@property
def mesh_as_arrays(self) -> dict:
## TODO: Cached
# Ensure Updated Geometry
log.debug('Updating View Layer')
bpy.context.view_layer.update()
# Compute Evaluted + Triangulated Mesh
log.debug('Casting BMesh of "%s" to Temporary Mesh', self.name)
_mesh = bpy.data.meshes.new(name='TemporaryMesh')
with self.mesh_as_bmesh(evaluate=True, triangulate=True) as bmesh_mesh:
bmesh_mesh.to_mesh(_mesh)
# Optimized Vertex Copy
## See <https://blog.michelanders.nl/2016/02/copying-vertices-to-numpy-arrays-in_4.html>
log.debug('Copying Vertices from "%s"', self.name)
verts = np.zeros(3 * len(_mesh.vertices), dtype=np.float64)
_mesh.vertices.foreach_get('co', verts)
verts.shape = (-1, 3)
# Optimized Triangle Copy
## To understand, read it, **carefully**.
log.debug('Copying Faces from "%s"', self.name)
faces = np.zeros(3 * len(_mesh.polygons), dtype=np.uint64)
_mesh.polygons.foreach_get('vertices', faces)
faces.shape = (-1, 3)
# Remove Temporary Mesh
log.debug('Removing Temporary Mesh')
bpy.data.meshes.remove(_mesh)
return {
'verts': verts,
'faces': faces,
}

View File

@ -0,0 +1,192 @@
import typing as typ
import bpy
import typing_extensions as typx
from ....utils import analyze_geonodes
from ....utils import logger
from .. import contracts as ct
log = logger.get(__name__)
ModifierType: typ.TypeAlias = typx.Literal['NODES', 'ARRAY']
NodeTreeInterfaceID: typ.TypeAlias = str
class ModifierAttrsNODES(typ.TypedDict):
node_group: bpy.types.GeometryNodeTree
inputs: dict[NodeTreeInterfaceID, typ.Any]
class ModifierAttrsARRAY(typ.TypedDict):
pass
ModifierAttrs: typ.TypeAlias = ModifierAttrsNODES | ModifierAttrsARRAY
MODIFIER_NAMES = {
'NODES': 'BLMaxwell_GeoNodes',
'ARRAY': 'BLMaxwell_Array',
}
####################
# - Read/Write Modifier Attributes
####################
def read_modifier(bl_modifier: bpy.types.Modifier) -> ModifierAttrs:
if bl_modifier.type == 'NODES':
return {
'node_group': bl_modifier.node_group,
}
elif bl_modifier.type == 'ARRAY':
raise NotImplementedError
raise NotImplementedError
def write_modifier(
bl_modifier: bpy.types.Modifier,
modifier_attrs: ModifierAttrs,
) -> bool:
"""Writes modifier attributes to the modifier, changing only what's needed.
Returns:
True if the modifier was altered.
"""
modifier_altered = False
if bl_modifier.type == 'NODES':
# Alter GeoNodes Group
if bl_modifier.node_group != modifier_attrs['node_group']:
log.info(
'Changing GeoNodes Modifier NodeTree from "%s" to "%s"',
str(bl_modifier.node_group),
str(modifier_attrs['node_group']),
)
bl_modifier.node_group = modifier_attrs['node_group']
modifier_altered = True
# Alter GeoNodes Input (Interface) Socket Values
## The modifier's dict-like setter actually sets NodeTree interface vals
## By setting the interface value, this particular NodeTree will change
geonodes_interface = analyze_geonodes.interface(
bl_modifier.node_group, direct='INPUT'
)
for (
socket_name,
raw_value,
) in modifier_attrs['inputs'].items():
iface_id = geonodes_interface[socket_name].identifier
# Alter Interface Value
if bl_modifier[iface_id] != raw_value:
# Determine IDPropertyArray Equality
## The equality above doesn't work for IDPropertyArrays.
## BUT, IDPropertyArrays must have a 'to_list' method.
## To do the comparison, we tuple-ify the IDPropertyArray.
## raw_value is always a tuple if it's listy.
if (
hasattr(bl_modifier[iface_id], 'to_list')
and tuple(bl_modifier[iface_id].to_list()) == raw_value
):
continue
# Determine int/float Mismatch
## Blender is strict; only floats can set float vals.
## We are less strict; if the user passes an int, that's okay.
if isinstance(
bl_modifier[iface_id],
float,
) and isinstance(raw_value, int):
value = float(raw_value)
bl_modifier[iface_id] = value
modifier_altered = True
## TODO: Altering existing values is much better for performance.
## - GC churn is real!
## - Especially since this is in a hot path
elif bl_modifier.type == 'ARRAY':
raise NotImplementedError
else:
raise NotImplementedError
return modifier_altered
####################
# - ManagedObj
####################
class ManagedBLModifier(ct.schemas.ManagedObj):
managed_obj_type = ct.ManagedObjType.ManagedBLModifier
_modifier_name: str | None = None
####################
# - BL Object Name
####################
@property
def name(self):
return self._modifier_name
@name.setter
def name(self, value: str) -> None:
## TODO: Handle name conflict within same BLObject
log.info(
'Changing BLModifier w/Name "%s" to Name "%s"', self._bl_object_name, value
)
self._modifier_name = value
####################
# - Allocation
####################
def __init__(self, name: str):
self.name = name
####################
# - Deallocation
####################
def free(self):
log.info('Freeing BLModifier w/Name "%s" (NOT IMPLEMENTED)', self.name)
## TODO: Implement
####################
# - Modifiers
####################
def bl_modifier(
self,
bl_object: bpy.types.Object,
modifier_type: ModifierType,
modifier_attrs: ModifierAttrs,
):
"""Creates a new modifier for the current `bl_object`.
- Modifier Type Names: <https://docs.blender.org/api/current/bpy_types_enum_items/object_modifier_type_items.html#rna-enum-object-modifier-type-items>
"""
# Remove Mismatching Modifier
if (
bl_modifier := bl_object.modifiers.get(self.name)
) and bl_modifier.type != modifier_type:
log.info(
'Removing (recreating) BLModifier "%s" on BLObject "%s" (existing modifier_type is "%s", but "%s" is requested)',
bl_modifier.name,
bl_object.name,
bl_modifier.type,
modifier_type,
)
self.free()
# Create Modifier
if not (bl_modifier := bl_object.modifiers.get(self.name)):
log.info(
'Creating BLModifier "%s" on BLObject "%s" with modifier_type "%s"',
self.name,
bl_object.name,
modifier_type,
)
bl_modifier = bl_object.modifiers.new(
name=self.name,
type=modifier_type,
)
if modifier_altered := write_modifier(bl_modifier, modifier_attrs):
bl_object.data.update()
return bl_modifier

View File

@ -7,6 +7,7 @@ import typing_extensions as typx
from ....utils import logger
from .. import contracts as ct
from .managed_bl_collection import managed_collection, preview_collection
log = logger.get(__name__)
@ -15,33 +16,6 @@ MODIFIER_NAMES = {
'NODES': 'BLMaxwell_GeoNodes',
'ARRAY': 'BLMaxwell_Array',
}
MANAGED_COLLECTION_NAME = 'BLMaxwell'
PREVIEW_COLLECTION_NAME = 'BLMaxwell Visible'
####################
# - BLCollection
####################
def bl_collection(
collection_name: str, view_layer_exclude: bool
) -> bpy.types.Collection:
# Init the "Managed Collection"
# Ensure Collection exists (and is in the Scene collection)
if collection_name not in bpy.data.collections:
collection = bpy.data.collections.new(collection_name)
bpy.context.scene.collection.children.link(collection)
else:
collection = bpy.data.collections[collection_name]
## Ensure synced View Layer exclusion
if (
layer_collection := bpy.context.view_layer.layer_collection.children[
collection_name
]
).exclude != view_layer_exclude:
layer_collection.exclude = view_layer_exclude
return collection
####################
@ -151,16 +125,9 @@ class ManagedBLObject(ct.schemas.ManagedObj):
If it's already included, do nothing.
"""
bl_object = self.bl_object(kind)
if (
bl_object.name
not in (
preview_collection := bl_collection(
PREVIEW_COLLECTION_NAME, view_layer_exclude=False
)
).objects
):
if bl_object.name not in preview_collection().objects:
log.info('Moving "%s" to Preview Collection', bl_object.name)
preview_collection.objects.link(bl_object)
preview_collection().objects.link(bl_object)
# Display Parameters
if kind == 'EMPTY' and empty_display_type is not None:
@ -180,14 +147,7 @@ class ManagedBLObject(ct.schemas.ManagedObj):
If it's already removed, do nothing.
"""
bl_object = self.bl_object(kind)
if (
bl_object.name
not in (
preview_collection := bl_collection(
PREVIEW_COLLECTION_NAME, view_layer_exclude=False
)
).objects
):
if bl_object.name not in preview_collection().objects:
log.info('Removing "%s" from Preview Collection', bl_object.name)
preview_collection.objects.unlink(bl_object)
@ -228,7 +188,7 @@ class ManagedBLObject(ct.schemas.ManagedObj):
if not (bl_object := bpy.data.objects.get(self.name)):
log.info(
'Creating "%s" with kind "%s"',
bl_object.name,
self.name,
kind,
)
if kind == 'MESH':
@ -246,9 +206,7 @@ class ManagedBLObject(ct.schemas.ManagedObj):
'Linking "%s" to Base Collection',
bl_object.name,
)
bl_collection(
MANAGED_COLLECTION_NAME, view_layer_exclude=True
).objects.link(bl_object)
managed_collection().objects.link(bl_object)
return bl_object

View File

@ -1,21 +1,23 @@
import tidy3d as td
import sympy as sp
import sympy.physics.units as spu
import typing as typ
import bpy
import sympy as sp
import sympy.physics.units as spu
import tidy3d as td
from ......utils import analyze_geonodes
from .... import contracts as ct
from .... import sockets
from .... import managed_objs
from .....assets.import_geonodes import import_geonodes
from .... import managed_objs, sockets
from ... import base
GEONODES_STRUCTURE_BOX = 'structure_box'
GEONODES_BOX = 'box'
class BoxStructureNode(base.MaxwellSimNode):
node_type = ct.NodeType.BoxStructure
bl_label = 'Box Structure'
use_sim_node_name = True
####################
# - Sockets
@ -27,15 +29,19 @@ class BoxStructureNode(base.MaxwellSimNode):
default_value=sp.Matrix([500, 500, 500]) * spu.nm
),
}
output_sockets = {
output_sockets: typ.ClassVar = {
'Structure': sockets.MaxwellStructureSocketDef(),
}
managed_obj_defs = {
'structure_box': ct.schemas.ManagedObjDef(
mk=lambda name: managed_objs.ManagedBLObject(name),
managed_obj_defs: typ.ClassVar = {
'mesh': ct.schemas.ManagedObjDef(
mk=lambda name: managed_objs.ManagedBLMesh(name),
name_prefix='',
)
),
'box': ct.schemas.ManagedObjDef(
mk=lambda name: managed_objs.ManagedBLModifier(name),
name_prefix='',
),
}
####################
@ -45,13 +51,15 @@ class BoxStructureNode(base.MaxwellSimNode):
'Structure',
input_sockets={'Medium', 'Center', 'Size'},
)
def compute_simulation(self, input_sockets: dict) -> td.Box:
def compute_structure(self, input_sockets: dict) -> td.Box:
medium = input_sockets['Medium']
_center = input_sockets['Center']
_size = input_sockets['Size']
center = as_unit_system(input_sockets['Center'], 'tidy3d')
size = as_unit_system(input_sockets['Size'], 'tidy3d')
#_center = input_sockets['Center']
#_size = input_sockets['Size']
center = tuple(spu.convert_to(_center, spu.um) / spu.um)
size = tuple(spu.convert_to(_size, spu.um) / spu.um)
#center = tuple(spu.convert_to(_center, spu.um) / spu.um)
#size = tuple(spu.convert_to(_size, spu.um) / spu.um)
return td.Structure(
geometry=td.Box(
@ -62,61 +70,43 @@ class BoxStructureNode(base.MaxwellSimNode):
)
####################
# - Preview - Changes to Input Sockets
# - Events
####################
@base.on_value_changed(
socket_name={'Center', 'Size'},
input_sockets={'Center', 'Size'},
managed_objs={'structure_box'},
managed_objs={'mesh', 'box'},
)
def on_value_changed__center_size(
self,
input_sockets: dict,
managed_objs: dict[str, ct.schemas.ManagedObj],
):
_center = input_sockets['Center']
center = tuple(
[float(el) for el in spu.convert_to(_center, spu.um) / spu.um]
)
center = as_unit_system(input_sockets['Center'], 'blender')
#center = tuple([float(el) for el in spu.convert_to(_center, spu.um) / spu.um])
## TODO: Implement + aggressively memoize as_unit_system
## - This should also understand that ex. Blender likes tuples, Tidy3D might like something else.
_size = input_sockets['Size']
size = tuple(
[float(el) for el in spu.convert_to(_size, spu.um) / spu.um]
)
## TODO: Preview unit system?? Presume um for now
size = as_unit_system(input_sockets['Size'], 'blender')
#size = tuple([float(el) for el in spu.convert_to(_size, spu.um) / spu.um])
# Retrieve Hard-Coded GeoNodes and Analyze Input
geo_nodes = bpy.data.node_groups[GEONODES_STRUCTURE_BOX]
geonodes_interface = analyze_geonodes.interface(
geo_nodes, direc='INPUT'
)
# Sync Modifier Inputs
managed_objs['structure_box'].sync_geonodes_modifier(
geonodes_node_group=geo_nodes,
geonodes_identifier_to_value={
geonodes_interface['Size'].identifier: size,
## TODO: Use 'bl_socket_map.value_to_bl`!
## - This accounts for auto-conversion, unit systems, etc. .
## - We could keep it in the node base class...
## - ...But it needs aligning with Blender, too. Hmm.
# Sync Attributes
managed_objs['mesh'].bl_object().location = center
managed_objs['box'].bl_modifier(managed_objs['mesh'].bl_object(), 'NODES', {
'node_group': import_geonodes(GEONODES_BOX, 'link'),
'inputs': {
'Size': size,
},
)
})
# Sync Object Position
managed_objs['structure_box'].bl_object('MESH').location = center
####################
# - Preview - Show Preview
####################
@base.on_show_preview(
managed_objs={'structure_box'},
managed_objs={'mesh'},
)
def on_show_preview(
self,
managed_objs: dict[str, ct.schemas.ManagedObj],
):
managed_objs['structure_box'].show_preview('MESH')
managed_objs['mesh'].show_preview()
self.on_value_changed__center_size()
@ -127,7 +117,5 @@ BL_REGISTER = [
BoxStructureNode,
]
BL_NODES = {
ct.NodeType.BoxStructure: (
ct.NodeCategory.MAXWELLSIM_STRUCTURES_PRIMITIVES
)
ct.NodeType.BoxStructure: (ct.NodeCategory.MAXWELLSIM_STRUCTURES_PRIMITIVES)
}