Skip to content

Instantly share code, notes, and snippets.

@PardhavMaradani

PardhavMaradani/README.md

Last active January 22, 2025 04:25
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save PardhavMaradani/3efd3b96188a97ff577d3151dfc5fd62 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save PardhavMaradani/3efd3b96188a97ff577d3151dfc5fd62 to your computer and use it in GitHub Desktop.
Download ZIP
PoC of MolecularNodes API
Raw
README.md

This is a quick PoC of a MolecularNodes API that allows different styles based on MDAnalysis selection strings using a single Blender object that represents a universe. This is to support the discussion in issue #5 and issue #11 of ggmolvis.

  • Copy / Install the mn_api.py file to the Blender's script/addons directory or execute from Blender's Text Editor
  • In Blender's Python Console run to verify this module can be loaded
import mn_api

Here are the API calls that match the ggmolvis example:

mn_api.load_example()
vu = mn_api.get_mn_universe('u')
vu.clear_styles()
vu.add_style('protein', style='cartoon')
vu.add_style('protein', style='surface', material='transparent')
vu.add_style('resid 127 40', style='ball_and_stick')
vu.clear_styles()
vu.add_style('all', 'ribbon', color=(0,1,1,0.1)) # rgba
Raw
mn_api.py
import bpy
import molecularnodes as mn
import os
MN_ADDON_DIR = os.path.dirname(mn.__file__)
MN_DATA_FILE = os.path.join(MN_ADDON_DIR, "assets", "MN_data_file_4.2.blend")
class MNUniverse:
def __init__(self, obj: str):
self.obj = bpy.data.objects[obj]
self.node_group = self.obj.modifiers["MolecularNodes"].node_group
def clear_styles(self):
self.obj.mn_trajectory_selections.clear()
self.node_group.nodes.clear()
_mn_default_node_group(self.node_group)
def add_style(
self,
selection: str,
style: str = "spheres",
material: str = None,
color: tuple = None,
):
self.obj.mn_trajectory_selections.add()
idx = len(self.obj.mn_trajectory_selections) - 1
selection_attr = "api_sel_" + str(idx)
self.obj.mn_trajectory_selections[idx].name = selection_attr
self.obj.mn_trajectory_selections[idx].selection_str = selection
_mn_add_style_frame(
self.node_group, selection_attr, style=style, material=material, color=color
)
def get_mn_universe(obj: str):
if obj not in bpy.data.objects:
print("Invalid Blender object name")
return None
return MNUniverse(obj)
def load_example():
import MDAnalysis as mda
from MDAnalysis.tests.datafiles import PSF, DCD
import warnings
warnings.filterwarnings("ignore")
if "u" in bpy.data.objects:
bpy.data.objects.remove(bpy.data.objects["u"])
if "MN_u" in bpy.data.node_groups:
bpy.data.node_groups.remove(bpy.data.node_groups["MN_u"])
u = mda.Universe(PSF, DCD)
# TODO: Below doesn't animate with frame changes - need to debug later
# t = mn.entities.Trajectory(universe=u)
# t.create_object(name="u", style="cartoon")
# Hence, simulate creating object through UI
bpy.context.scene.MN_import_md_topology = PSF
bpy.context.scene.MN_import_md_trajectory = DCD
bpy.context.scene.MN_import_md_name = "u"
bpy.context.scene.mn.import_style = "spheres"
bpy.ops.mn.import_trajectory()
bpy.context.scene.frame_start = 0
bpy.context.scene.frame_end = u.trajectory.n_frames - 1
_import_all_mn_nodes()
_import_all_mn_materials()
_fix_mn_default_material_alpha()
def _fix_mn_default_material_alpha():
# node group: MN Color Input
mn_color_input = bpy.data.node_groups["MN Color Input"]
# node: Attribute
attribute = mn_color_input.nodes["Attribute"]
# node: Group Output
group_output = mn_color_input.nodes["Group Output"]
# Socket Alpha
alpha_socket = mn_color_input.interface.new_socket(
name="Alpha", in_out="OUTPUT", socket_type="NodeSocketFloat"
)
alpha_socket.default_value = 1.0
alpha_socket.min_value = 0
alpha_socket.max_value = 1
alpha_socket.subtype = "NONE"
alpha_socket.attribute_domain = "POINT"
# attribute.Alpha -> group_output.Alpha
mn_color_input.links.new(attribute.outputs["Alpha"], group_output.inputs["Alpha"])
# material: MN Default
mn_default = bpy.data.materials["MN Default"].node_tree
# node: Group
group = mn_default.nodes["Group"]
# node: Principled BSDF
principled_bsdf = mn_default.nodes["Principled BSDF"]
# group.Alpha -> principled_bsdf.Alpha
mn_default.links.new(group.outputs["Alpha"], principled_bsdf.inputs["Alpha"])
def _mn_default_node_group(node_group: bpy.types.GeometryNodeTree):
# node: Group Input
group_input = node_group.nodes.new("NodeGroupInput")
group_input.name = "Group Input"
# node: Group Output
group_output = node_group.nodes.new("NodeGroupOutput")
group_output.name = "Group Output"
group_output.is_active_output = True
# node: Color Common
color_common = node_group.nodes.new("GeometryNodeGroup")
color_common.name = "Color Common"
color_common.node_tree = bpy.data.node_groups["Color Common"]
# node: Color Attribute Random
color_attribute_random = node_group.nodes.new("GeometryNodeGroup")
color_attribute_random.name = "Color Attribute Random"
color_attribute_random.node_tree = bpy.data.node_groups["Color Attribute Random"]
# node: Join Geometry
join_geometry = node_group.nodes.new("GeometryNodeJoinGeometry")
join_geometry.name = "Join Geometry"
# Set locations
group_input.location = (0.0, 0.0)
group_output.location = (880.0, 0.0)
color_common.location = (-50.0, -150.0)
color_attribute_random.location = (-300.0, -150.0)
join_geometry.location = (700.0, -5.0)
# Set dimensions
group_input.width, group_input.height = 140.0, 100.0
group_output.width, group_output.height = 140.0, 100.0
color_common.width, color_common.height = 180.0, 100.0
color_attribute_random.width, color_attribute_random.height = 180.0, 100.0
# initialize node_group links
node_group.links.new(
color_attribute_random.outputs["Color"], color_common.inputs["Carbon"]
)
node_group.links.new(
join_geometry.outputs["Geometry"], group_output.inputs["Geometry"]
)
return node_group
def _mn_add_style_frame(
node_group: bpy.types.GeometryNodeTree,
selection_attr: str,
style: str = "spheres",
material: str = None,
color: tuple = None,
):
# node: Named Attribute
named_attribute = node_group.nodes.new("GeometryNodeInputNamedAttribute")
named_attribute.name = "Named Attribute"
named_attribute.data_type = "BOOLEAN"
named_attribute.inputs["Name"].default_value = selection_attr
# node: Set Color
set_color = node_group.nodes.new("GeometryNodeGroup")
set_color.label = "Set Color"
set_color.name = "Set Color"
set_color.node_tree = bpy.data.node_groups["Set Color"]
# node: Style
style = _mn_get_style_node(node_group, style=style, material=material)
# node: Frame
frame = node_group.nodes.new("NodeFrame")
frame.label = selection_attr
frame.name = "Frame"
frame.label_size = 20
frame.shrink = True
# Set parents
named_attribute.parent = frame
set_color.parent = frame
style.parent = frame
# Set locations
named_attribute.location = (-115.0, -55.0)
set_color.location = (-117.0, 96.0)
style.location = (86.0, 94.0)
frame.location = (353.0, -172.0)
# Set dimensions
named_attribute.width, named_attribute.height = 180.0, 100.0
set_color.width, set_color.height = 180.0, 100.0
style.width, style.height = 180.0, 100.0
frame.width, frame.height = 443.0, 340.0
# initialize node_group links
node_group.links.new(set_color.outputs["Atoms"], style.inputs["Atoms"])
node_group.links.new(
named_attribute.outputs["Attribute"], style.inputs["Selection"]
)
node_group.links.new(
node_group.nodes["Group Input"].outputs["Geometry"], set_color.inputs["Atoms"]
)
if color is None:
node_group.links.new(
node_group.nodes["Color Common"].outputs["Color"], set_color.inputs["Color"]
)
else:
set_color.inputs["Color"].default_value = color
node_group.links.new(
style.outputs["Geometry"], node_group.nodes["Join Geometry"].inputs["Geometry"]
)
return node_group
def _mn_get_style_node(
node_group: bpy.types.GeometryNodeTree,
style: str = "spheres",
material: str = None,
):
materials_mapping = {
"default": "MN Default",
"flat": "MN Flat Outline",
"squishy": "MN Squishy",
"transparent": "MN Transparent Outline",
"ambient": "MN Ambient Occlusion",
}
# node: Style
node = node_group.nodes.new("GeometryNodeGroup")
node.label = "Style"
node.name = "Style"
if style == "ball_and_stick":
node.node_tree = bpy.data.node_groups["Style Ball and Stick"]
elif style == "cartoon":
node.node_tree = bpy.data.node_groups["Style Cartoon"]
node.inputs["DSSP"].default_value = True
elif style == "ribbon":
node.node_tree = bpy.data.node_groups["Style Ribbon"]
elif style == "spheres":
node.node_tree = bpy.data.node_groups["Style Spheres"]
node.inputs["Sphere As Mesh"].default_value = True
elif style == "sticks":
node.node_tree = bpy.data.node_groups["Style Sticks"]
elif style == "surface":
node.node_tree = bpy.data.node_groups["Style Surface"]
if material in materials_mapping:
node.inputs["Material"].default_value = bpy.data.materials[
materials_mapping[material]
]
else:
node.inputs["Material"].default_value = bpy.data.materials["MN Default"]
return node
def _import_all_mn_nodes():
with bpy.data.libraries.load(MN_DATA_FILE, link=False) as (data_from, data_to):
data_to.node_groups = data_from.node_groups
def _import_mn_node(node_name: str):
with bpy.data.libraries.load(MN_DATA_FILE, link=False) as (data_from, data_to):
data_to.node_groups.append(node_name)
def _import_all_mn_materials():
with bpy.data.libraries.load(MN_DATA_FILE, link=False) as (data_from, data_to):
data_to.materials = data_from.materials
def _import_mn_material(material_name: str):
with bpy.data.libraries.load(MN_DATA_FILE, link=False) as (data_from, data_to):
data_to.materials.append(material_name)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment