Skip to content

Instantly share code, notes, and snippets.

@PardhavMaradani

PardhavMaradani/README.md

Last active January 14, 2025 02:23
Show Gist options

  • Select an option

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

Select an option

Learn more about clone URLs
Save PardhavMaradani/13d77d8c107e3d3c17228dc8598b11bd to your computer and use it in GitHub Desktop.
Download ZIP
PoC of Basic Annotations, Density and Pore visualization of MDAnalysis results in Blender
Raw
README.md

This gist has instructions and code for the Visualization of MDAnalysis results in Blender discussion.

Basic Annotations

  • Copy / Install the bmda.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 bmda as b
  • The basic annotations API is as follows:
    • b.show_annotations() to show any annotations added
    • b.hide_annotations() to hide any annotations added
    • b.clear_annotations() to clear all added annotations
    • b.show_atom_info(...) to annotate atom details
      • Eg: b.show_atom_info('u', 'index 1139') shows the atom name at index 1139 in the universe
        • The first parameter (here u) is the object name (string) of the trajectory that was loaded through MolecularNodes
        • The second parameter is a string that is used to get the AtomGroup selection using .select_atoms of MDAnalysis universe
    • b.show_bond_angle(...) to show bond angle details between 3 Atoms
      • Eg: b.show_bond_angle('u', ['index 1143', 'index 1142', 'index 1144'])
        • The selections can also be an array of strings to ensure that the correct order is obtained in the AtomGroup as above, especially for angles
    • b.show_com(...) to show the center of mass of a selection
      • Eg: b.show_com('u', 'protein')
    • b.show_com_distance(...) to show the distance between two center of masses of two selections
      • Eg: b.show_com_distance('u', 'resid 1 and name CA', 'resid 129 and name CA')
    • b.show_canonical_dihedrals(...) to show the canonical dihedral angles in a given residue
      • Eg: b.show_canonical_dihedrals('u', 200) shows the phi, psi, omega and chi1 angles of residue 200

All the above examples are for the trajectory mda.Universe(PSF, DCD) and tested only for Molecular Nodes version 4.2.10 in Blender 4.3.2

There are a couple of simple utility API calls:

  • b.get_universe(...) to get the MDAnalysis universe associated with a Blender object
    • Eg: b.get_universe('u') returns the universe for object named u
  • b.print_annotations() to print the current annotation list

An optional options parameter can be passed to most of the API above that can set the font size, font color, line color, text displayed etc.

Eg:

b.show_atom_info('u', 'index 1139', {'showRes': True, 'showSeg': True, 'fontColor': (1, 0, 0, 1), 'fontSize': 12})

shows the residue name and the segment name along with the atom details

b.show_com_distance('u', ['resid 1 and name CA'], 'resid 129 and name CA', {'text1': 'Res 1|COM', 'text2': 'Res 129|COM', 'fontColor': (1, 0, 0, 1), 'fontSize': 12, 'lineColor': (0, 1, 0, 1)})

shows the COM distance with a custom text, font, color, size etc

Density Visualization

  • Copy / Install the mda_density_viz.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 mda_density_viz as d

If your Blender's python environment has all the standard MDAnalysis packages installed, you can just run:

d.load_example()

to load the analysis results from Calculating the solvent density around a protein section of the user guide. Please note that this will take several seconds to finish loading, but once done, you should see something similar to what is shown in the video where you can change the modifier values to visualize different iso levels, change colors/alpha, slice, see contours etc.

If you want to load other examples of your own:

  • A couple of other useful methods available are:
    • d.dx2vdb(...) to convert any dx file to a vdb file
    • d.load_vdb(...) to load any vdb file - this will also setup the required geometry nodes for visualization

All the geometry node setup happens via the generated code from Node To Python Blender Extension.

Pore Visualization

  • Copy / Install the mda_pore_viz.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 mda_pore_viz as p

To generate the pore center line:

p.draw_pore_center_line("hole1.sph")

where hole1.sph is the .sph generated by HOLE2 analysis

To generate the pore:

p.create_pore("hole1.sph")

Note that the above command can take several seconds due the the poor performance of the boolean modifier used currently.

Raw
bmda.py
import bpy
import blf # type: ignore
import gpu # type: ignore
from gpu_extras.batch import batch_for_shader # type: ignore
from bpy_extras import view3d_utils # type: ignore
from math import fabs, radians, sqrt, cos, sin
# External
def show_annotations():
global _draw_handler
_draw_handler = bpy.types.SpaceView3D.draw_handler_add(
_draw_annotations_handler, (None, bpy.context), "WINDOW", "POST_PIXEL"
)
_redraw(bpy.context)
def hide_annotations():
global _draw_handler
if _draw_handler is not None:
bpy.types.SpaceView3D.draw_handler_remove(_draw_handler, "WINDOW")
_redraw(bpy.context)
_draw_handler = None
def clear_annotations():
global _annotations
_annotations = []
_redraw(bpy.context)
def print_annotations():
print(_annotations)
def show_atom_info(obj_name, selection, options=None):
if not _valid_object(obj_name):
print("Invalid Object", obj_name)
return
selection_array = _make_selection_array(selection)
ag = _get_selection(obj_name, selection_array)
if ag is None:
print("Invalid selection", obj_name, selection)
return
print(repr(ag))
annotation = {
"type": "atom_info",
"obj_name": obj_name,
"selection": selection_array,
"options": options,
}
_annotations.append(annotation)
_redraw(bpy.context)
def show_bond_angle(obj_name, selection, options=None):
if not _valid_object(obj_name):
print("Invalid Object", obj_name)
return
selection_array = _make_selection_array(selection)
ag = _get_selection(obj_name, selection_array)
if ag is None:
print("Invalid selection", obj_name, selection)
return
if ag.n_atoms != 3:
print("Need 3 atoms, selection returned", ag.n_atoms)
return
print(repr(ag))
annotation = {
"type": "bond_angle",
"obj_name": obj_name,
"selection": selection_array,
"options": options,
}
_annotations.append(annotation)
_redraw(bpy.context)
def show_com(obj_name, selection, options=None):
if not _valid_object(obj_name):
print("Invalid Object", obj_name)
return
selection_array = _make_selection_array(selection)
ag = _get_selection(obj_name, selection_array)
if ag is None:
print("Invalid selection", obj_name, selection)
return
print(repr(ag))
annotation = {
"type": "com",
"obj_name": obj_name,
"selection": selection_array,
"options": options,
}
_annotations.append(annotation)
_redraw(bpy.context)
def show_com_distance(obj_name, selection1, selection2, options=None):
if not _valid_object(obj_name):
print("Invalid Object", obj_name)
return
selection1_array = _make_selection_array(selection1)
ag1 = _get_selection(obj_name, selection1_array)
if ag1 is None:
print("Invalid selection", obj_name, selection1)
return
selection2_array = _make_selection_array(selection2)
ag2 = _get_selection(obj_name, selection2_array)
if ag1 is None:
print("Invalid selection", obj_name, selection2)
return
print(repr(ag1), repr(ag2))
annotation = {
"type": "com_distance",
"obj_name": obj_name,
"selection1": selection1_array,
"selection2": selection2_array,
"options": options,
}
_annotations.append(annotation)
_redraw(bpy.context)
def show_canonical_dihedrals(obj_name, resid, options=None):
if not _valid_object(obj_name):
print("Invalid Object", obj_name)
return
u = _get_universe(obj_name)
try:
r = u.residues[resid - 1]
print(repr(r))
phi = r.phi_selection()
print("phi", phi.names, phi.dihedral.value())
psi = r.psi_selection()
print("psi", psi.names, psi.dihedral.value())
omega = r.omega_selection()
print("omega", omega.names, omega.dihedral.value())
chi1 = r.chi1_selection()
print("chi1", chi1.names, chi1.dihedral.value())
except Exception as e:
print("Exception :", e)
return
annotation = {
"type": "dihedrals",
"obj_name": obj_name,
"resid": resid,
"options": options,
}
_annotations.append(annotation)
_redraw(bpy.context)
def get_universe(obj_name):
return _get_universe(obj_name)
# Internal
_fmt = "%1.2f"
_annotations = []
_viewport = (0, 0)
_draw_handler = None
_rad45 = radians(45)
_rad315 = radians(315)
_bsf = 0.01
_default_color = (1, 1, 1, 1)
_default_font_size = 16
_default_line_width = 1.0
_default_arrow_size = 16
_default_text_align = "C"
_default_text_rotation = 0.0
_shader_line = (
gpu.shader.from_builtin("POLYLINE_UNIFORM_COLOR")
if not bpy.app.background
else None
)
def _set_viewport():
global _viewport
_viewport = (bpy.context.region.width, bpy.context.region.height)
def _valid_object(obj_name):
if obj_name in bpy.context.scene.objects and hasattr(
bpy.context.scene.objects[obj_name], "uuid"
):
uuid = bpy.context.scene.objects[obj_name].uuid
if (
hasattr(bpy.context.scene, "MNSession")
and uuid in bpy.context.scene.MNSession.trajectories
):
return True
return False
def _make_selection_array(selection):
selection_array = selection
if type(selection) is str:
selection_array = [selection]
return selection_array
def _get_selection(obj_name, selection):
u = _get_universe(obj_name)
try:
ag = u.select_atoms(*selection)
return ag
except Exception as e:
print("Exception :", e)
return None
def _redraw(context):
if _draw_handler is not None:
area = next(area for area in context.screen.areas if area.type == "VIEW_3D")
area.tag_redraw()
def _draw_annotations_handler(self, context):
_set_viewport()
region = bpy.context.region
if not context.space_data.region_quadviews:
rv3d = bpy.context.space_data.region_3d
else:
if context.area.type != "VIEW_3D" or context.space_data.type != "VIEW_3D":
return
i = -1
for region in context.area.regions:
if region.type == "WINDOW":
i += 1
if context.region.id == region.id:
break
else:
return
rv3d = context.space_data.region_quadviews[i]
for annotation in _annotations:
if annotation["type"] == "atom_info":
_draw_atom_info(region, rv3d, annotation)
elif annotation["type"] == "bond_angle":
_draw_bond_angle(region, rv3d, annotation)
elif annotation["type"] == "com":
_draw_com(region, rv3d, annotation)
elif annotation["type"] == "com_distance":
_draw_com_distance(region, rv3d, annotation)
elif annotation["type"] == "dihedrals":
_draw_dihedrals(region, rv3d, annotation)
def _get_universe(obj_name):
uuid = bpy.context.scene.objects[obj_name].uuid
trajectory = bpy.context.scene.MNSession.trajectories[uuid]
return trajectory.universe
def _draw_atom_info(region, rv3d, annotation):
ag = _get_selection(annotation["obj_name"], annotation["selection"])
options = annotation["options"]
for a in ag:
text = a.name
if options is not None:
if options.get("showRes"):
text += "|res " + str(a.resid) + ", " + a.resname
if options.get("showSeg"):
text += "|seg " + a.segid
_draw_text_3d(region, rv3d, a.position * _bsf, text, options=options)
def _draw_bond_angle(region, rv3d, annotation):
ag = _get_selection(annotation["obj_name"], annotation["selection"])
options = annotation["options"]
for i in range(3):
a = ag.atoms[i]
text = a.name
_draw_text_3d(region, rv3d, a.position * _bsf, text, options=options)
v1 = ag.atoms[0].position * _bsf
v2 = ag.atoms[1].position * _bsf
v3 = ag.atoms[2].position * _bsf
dist = _distance(v2, v1)
fdist = (_fmt % (dist / _bsf)) + " Å"
v21o = _interpolate3d(v2, v1, dist * 0.2)
_draw_arrow_line_3d(
region,
rv3d,
v2,
v1,
text=fdist,
leftArrow=False,
rightArrow=True,
options=options,
)
dist = _distance(v2, v3)
fdist = (_fmt % (dist / _bsf)) + " Å"
_draw_arrow_line_3d(
region,
rv3d,
v2,
v3,
text=fdist,
leftArrow=False,
rightArrow=True,
options=options,
)
v23o = _interpolate3d(v2, v3, dist * 0.2)
angle = _fmt % ag.angle.value() + " °"
_draw_arrow_line_3d(region, rv3d, v21o, v23o, text=angle, options=options)
def _draw_com(region, rv3d, annotation):
ag = _get_selection(annotation["obj_name"], annotation["selection"])
com = ag.center_of_mass() * _bsf
text = ",".join(annotation["selection"]) + "|" + "COM"
options = annotation["options"]
if options is not None:
text = options.get("text") or text
_draw_text_3d(region, rv3d, com, text, options=options)
def _draw_com_distance(region, rv3d, annotation):
ag1 = _get_selection(annotation["obj_name"], annotation["selection1"])
com1 = ag1.center_of_mass() * _bsf
text = ",".join(annotation["selection1"]) + "|" + "COM"
options = annotation["options"]
if options is not None:
text = options.get("text1") or text
_draw_text_3d(region, rv3d, com1, text, options=options)
ag2 = _get_selection(annotation["obj_name"], annotation["selection2"])
com2 = ag2.center_of_mass() * _bsf
text = ",".join(annotation["selection2"]) + "|" + "COM"
if options is not None:
text = options.get("text2") or text
_draw_text_3d(region, rv3d, com2, text, options=options)
dist = _distance(com1, com2) / _bsf
fdist = ("%1.2f" % dist) + " Å"
_draw_arrow_line_3d(
region,
rv3d,
com1,
com2,
text=fdist,
leftArrow=True,
rightArrow=True,
options=options,
)
def _draw_dihedrals(region, rv3d, annotation):
u = _get_universe(annotation["obj_name"])
r = u.residues[annotation["resid"] - 1]
selections = [
r.phi_selection(),
r.psi_selection(),
r.omega_selection(),
r.chi1_selection(),
]
symbols = ["ϕ", "ψ", "ω", "χ1"]
options = annotation["options"]
for i in range(4):
for j in range(4):
a = selections[i].atoms[j]
_draw_text_3d(region, rv3d, a.position * _bsf, a.name, options=options)
if j != 3:
b = selections[i].atoms[j + 1]
text = None
if j == 1:
text = (
symbols[i]
+ " = "
+ (_fmt % selections[i].dihedral.value())
+ " °"
)
_draw_arrow_line_3d(
region,
rv3d,
a.position * _bsf,
b.position * _bsf,
text=text,
rightArrow=(j == 1),
options=options,
)
def _in_viewport(pos2d):
x, y = pos2d
vw, vh = _viewport
if x < 0 or x > vw or y < 0 or y > vh:
return False
return True
def _get_2d_point(region, rv3d, point3d):
return view3d_utils.location_3d_to_region_2d(region, rv3d, point3d)
def _draw_text_2d(pos2d, display_text, options=None):
if pos2d is None or not _in_viewport(pos2d):
return
rgba = _default_color
fsize = _default_font_size
align = _default_text_align
text_rot = _default_text_rotation
if options is not None:
rgba = options.get("fontColor") or _default_color
fsize = options.get("fontSize") or _default_font_size
align = options.get("align") or _default_text_align
text_rot = options.get("textRotation") or _default_text_rotation
gap = 12
x_pos, y_pos = pos2d
font_id = 0
ui_scale = bpy.context.preferences.system.ui_scale
blf.size(font_id, round(fsize * ui_scale))
# height of one line
mwidth, mheight = blf.dimensions(font_id, "Tp") # uses high/low letters
# split lines
mylines = display_text.split("|")
idx = len(mylines) - 1
# -------------------
# Draw all text lines
# -------------------
for line in mylines:
text_width, text_height = blf.dimensions(font_id, line)
if align == "C":
new_x = x_pos - text_width / 2
elif align == "R":
new_x = x_pos - text_width - gap
else:
new_x = x_pos
blf.enable(font_id, blf.ROTATION)
blf.rotation(font_id, text_rot)
# calculate new Y position
new_y = y_pos + (mheight * idx)
# Draw
blf.position(font_id, new_x, new_y, 0)
blf.color(font_id, rgba[0], rgba[1], rgba[2], rgba[3])
blf.draw(font_id, " " + line)
# sub line
idx -= 1
if align == "L":
blf.disable(font_id, blf.ROTATION)
def _draw_text_3d(region, rv3d, v, display_text, options=None):
pos2dv = _get_2d_point(region, rv3d, v)
_draw_text_2d(pos2dv, display_text, options)
def _distance(v1, v2):
d = sqrt((v2[0] - v1[0]) ** 2 + (v2[1] - v1[1]) ** 2 + (v2[2] - v1[2]) ** 2)
return d
def _interpolate3d(v1, v2, d1):
# calculate vector
v = (v2[0] - v1[0], v2[1] - v1[1], v2[2] - v1[2])
# calculate distance between points
d0 = _distance(v1, v2)
# calculate interpolate factor (distance from origin / distance total)
# if d1 > d0, the point is projected in 3D space
if d0 > 0:
x = d1 / d0
else:
x = d1
final = (v1[0] + (v[0] * x), v1[1] + (v[1] * x), v1[2] + (v[2] * x))
return final
def _draw_line_2d(v1, v2, options=None):
if v1 is None or v2 is None:
return
rgba = _default_color
lineWidth = _default_line_width
if options is not None:
rgba = options.get("lineColor") or _default_color
lineWidth = options.get("lineWidth") or _default_line_width
coords = [(v1[0], v1[1], 0), (v2[0], v2[1], 0)]
gpu.state.blend_set("ALPHA")
batch = batch_for_shader(_shader_line, "LINES", {"pos": coords})
try:
_shader_line.bind()
_shader_line.uniform_float("color", rgba)
_shader_line.uniform_float("lineWidth", lineWidth)
_shader_line.uniform_float("viewportSize", _viewport)
batch.draw(_shader_line)
except Exception as e:
print(e)
def _draw_arrow_line_2d(v1, v2, leftArrow=False, rightArrow=False, options=None):
if v1 is None or v2 is None:
return
arrowSize = _default_arrow_size
if options is not None:
arrowSize = options.get("arrowSize") or _default_arrow_size
_draw_line_2d(v1, v2, options=options)
if leftArrow:
v = _interpolate3d((v1[0], v1[1], 0.0), (v2[0], v2[1], 0.0), arrowSize)
v1i = (v[0] - v1[0], v[1] - v1[1])
v1a = (
int(v1i[0] * cos(_rad45) - v1i[1] * sin(_rad45) + v1[0]),
int(v1i[1] * cos(_rad45) + v1i[0] * sin(_rad45)) + v1[1],
)
v1b = (
int(v1i[0] * cos(_rad315) - v1i[1] * sin(_rad315) + v1[0]),
int(v1i[1] * cos(_rad315) + v1i[0] * sin(_rad315) + v1[1]),
)
_draw_line_2d(v1, v1a, options=options)
_draw_line_2d(v1, v1b, options=options)
if rightArrow:
v = _interpolate3d((v2[0], v2[1], 0.0), (v1[0], v1[1], 0.0), arrowSize)
v2i = (v[0] - v2[0], v[1] - v2[1])
v2a = (
int(v2i[0] * cos(_rad45) - v2i[1] * sin(_rad45) + v2[0]),
int(v2i[1] * cos(_rad45) + v2i[0] * sin(_rad45)) + v2[1],
)
v2b = (
int(v2i[0] * cos(_rad315) - v2i[1] * sin(_rad315) + v2[0]),
int(v2i[1] * cos(_rad315) + v2i[0] * sin(_rad315) + v2[1]),
)
_draw_line_2d(v2, v2a, options=options)
_draw_line_2d(v2, v2b, options=options)
def _draw_arrow_line_3d(
region, rv3d, v1, v2, text=None, leftArrow=False, rightArrow=False, options=None
):
if v1 is None or v2 is None:
return
pos2dv1 = _get_2d_point(region, rv3d, v1)
pos2dv2 = _get_2d_point(region, rv3d, v2)
_draw_arrow_line_2d(pos2dv1, pos2dv2, leftArrow, rightArrow, options=options)
if text is not None:
dist = _distance(v1, v2)
midpoint3d = _interpolate3d(v1, v2, fabs(dist / 2))
_draw_text_3d(region, rv3d, midpoint3d, text, options=options)
Raw
mda_density_viz.py
import bpy
"""
import mda_density_viz as d
d.load_example()
"""
def load_example(datadir="/tmp/"):
# datadir = bpy.app.tempdir
dxfile = datadir + "water.dx"
vdbfile = datadir + "water.vdb"
_run_density_analysis(dxfile)
dx2vdb(dxfile, vdbfile)
load_vdb(vdbfile)
def dx2vdb(infile, outfile):
import pyopenvdb as vdb
from gridData import Grid
import numpy as np
d = Grid(infile)
vgrid = vdb.FloatGrid()
vgrid.copyFromArray(d.grid)
vgrid.gridClass = vdb.GridClass.FOG_VOLUME
vgrid.name = "density"
vgrid.transform.translate(np.array(d.origin) * 0.01)
vgrid.transform.scale(np.array(d.delta) * 0.01)
vdb.write(outfile, grids=[vgrid])
def load_vdb(filepath):
bpy.ops.object.volume_import(filepath=filepath)
_setup_geometry_nodes()
# setup geometry node tree for active volume object
def _setup_geometry_nodes():
p_material = _p_material_node_group()
bpy.context.active_object.data.materials.append(p_material)
n_material = _n_material_node_group()
bpy.context.active_object.data.materials.append(n_material)
bpy.ops.object.modifier_add(type="NODES")
bpy.context.object.modifiers["GeometryNodes"].node_group = (
_mda_density_viz_node_group()
)
bpy.context.object.modifiers["GeometryNodes"].node_group.is_modifier = True
def _run_density_analysis(dxfile):
import MDAnalysis as mda
from MDAnalysis.tests.datafiles import TPR, XTC
from MDAnalysis.analysis import density
from MDAnalysis import transformations as trans
import molecularnodes as mn
u = mda.Universe(TPR, XTC)
protein = u.select_atoms("protein")
water = u.select_atoms("resname SOL")
workflow = [
trans.unwrap(u.atoms), # unwrap all fragments
trans.center_in_box(
protein, center="geometry" # move atoms so protein
), # is centered
trans.wrap(
water, compound="residues" # wrap water back into box
), # keep each water whole
trans.fit_rot_trans(
protein, protein, weights="mass" # align protein to first frame
),
]
u.trajectory.add_transformations(*workflow)
traj = mn.entities.trajectory.Trajectory(universe=u)
traj.create_object(name="u", style="cartoon")
ow = u.select_atoms("name OW")
dens = density.DensityAnalysis(ow, delta=4.0, padding=2)
dens.run()
d = dens.results.density
d.convert_density("TIP4P")
d.export(dxfile)
# auto-generated by 'Node To Python' extension
# https://extensions.blender.org/add-ons/node-to-python/
# initialize p_material node group
def _p_material_node_group():
mat = bpy.data.materials.new(name="p_material")
mat.use_nodes = True
p_material = mat.node_tree
# start with a clean node tree
for node in p_material.nodes:
p_material.nodes.remove(node)
p_material.color_tag = "NONE"
p_material.description = ""
p_material.default_group_node_width = 140
# p_material interface
# initialize p_material nodes
# node Material Output
material_output = p_material.nodes.new("ShaderNodeOutputMaterial")
material_output.name = "Material Output"
material_output.is_active_output = True
material_output.target = "ALL"
# Displacement
material_output.inputs[2].default_value = (0.0, 0.0, 0.0)
# Thickness
material_output.inputs[3].default_value = 0.0
# node Principled BSDF
principled_bsdf = p_material.nodes.new("ShaderNodeBsdfPrincipled")
principled_bsdf.name = "Principled BSDF"
principled_bsdf.distribution = "MULTI_GGX"
principled_bsdf.subsurface_method = "RANDOM_WALK"
# Metallic
principled_bsdf.inputs[1].default_value = 0.0
# Roughness
principled_bsdf.inputs[2].default_value = 0.5
# IOR
principled_bsdf.inputs[3].default_value = 1.5
# Normal
principled_bsdf.inputs[5].default_value = (0.0, 0.0, 0.0)
# Diffuse Roughness
principled_bsdf.inputs[7].default_value = 0.0
# Subsurface Weight
principled_bsdf.inputs[8].default_value = 0.0
# Subsurface Radius
principled_bsdf.inputs[9].default_value = (1.0, 0.2, 0.1)
# Subsurface Scale
principled_bsdf.inputs[10].default_value = 0.05
# Subsurface Anisotropy
principled_bsdf.inputs[12].default_value = 0.0
# Specular IOR Level
principled_bsdf.inputs[13].default_value = 0.5
# Specular Tint
principled_bsdf.inputs[14].default_value = (1.0, 1.0, 1.0, 1.0)
# Anisotropic
principled_bsdf.inputs[15].default_value = 0.0
# Anisotropic Rotation
principled_bsdf.inputs[16].default_value = 0.0
# Tangent
principled_bsdf.inputs[17].default_value = (0.0, 0.0, 0.0)
# Transmission Weight
principled_bsdf.inputs[18].default_value = 0.0
# Coat Weight
principled_bsdf.inputs[19].default_value = 0.0
# Coat Roughness
principled_bsdf.inputs[20].default_value = 0.03
# Coat IOR
principled_bsdf.inputs[21].default_value = 1.5
# Coat Tint
principled_bsdf.inputs[22].default_value = (1.0, 1.0, 1.0, 1.0)
# Coat Normal
principled_bsdf.inputs[23].default_value = (0.0, 0.0, 0.0)
# Sheen Weight
principled_bsdf.inputs[24].default_value = 0.0
# Sheen Roughness
principled_bsdf.inputs[25].default_value = 0.5
# Sheen Tint
principled_bsdf.inputs[26].default_value = (1.0, 1.0, 1.0, 1.0)
# Emission Color
principled_bsdf.inputs[27].default_value = (1.0, 1.0, 1.0, 1.0)
# Emission Strength
principled_bsdf.inputs[28].default_value = 0.0
# Thin Film Thickness
principled_bsdf.inputs[29].default_value = 0.0
# Thin Film IOR
principled_bsdf.inputs[30].default_value = 1.33
# node Attribute
attribute = p_material.nodes.new("ShaderNodeAttribute")
attribute.name = "Attribute"
attribute.attribute_name = "iso_positive"
attribute.attribute_type = "GEOMETRY"
# Set locations
material_output.location = (300.0, 300.0)
principled_bsdf.location = (32.0, 305.0)
attribute.location = (-150.0, 234.0)
# Set dimensions
material_output.width, material_output.height = 140.0, 100.0
principled_bsdf.width, principled_bsdf.height = 240.0, 100.0
attribute.width, attribute.height = 140.0, 100.0
# initialize p_material links
# principled_bsdf.BSDF -> material_output.Surface
p_material.links.new(principled_bsdf.outputs[0], material_output.inputs[0])
# attribute.Color -> principled_bsdf.Base Color
p_material.links.new(attribute.outputs[0], principled_bsdf.inputs[0])
# attribute.Alpha -> principled_bsdf.Alpha
p_material.links.new(attribute.outputs[3], principled_bsdf.inputs[4])
return mat
# auto-generated by 'Node To Python' extension
# https://extensions.blender.org/add-ons/node-to-python/
# initialize n_material node group
def _n_material_node_group():
mat = bpy.data.materials.new(name="n_material")
mat.use_nodes = True
n_material = mat.node_tree
# start with a clean node tree
for node in n_material.nodes:
n_material.nodes.remove(node)
n_material.color_tag = "NONE"
n_material.description = ""
n_material.default_group_node_width = 140
# n_material interface
# initialize n_material nodes
# node Material Output
material_output = n_material.nodes.new("ShaderNodeOutputMaterial")
material_output.name = "Material Output"
material_output.is_active_output = True
material_output.target = "ALL"
# Displacement
material_output.inputs[2].default_value = (0.0, 0.0, 0.0)
# Thickness
material_output.inputs[3].default_value = 0.0
# node Principled BSDF
principled_bsdf = n_material.nodes.new("ShaderNodeBsdfPrincipled")
principled_bsdf.name = "Principled BSDF"
principled_bsdf.distribution = "MULTI_GGX"
principled_bsdf.subsurface_method = "RANDOM_WALK"
# Metallic
principled_bsdf.inputs[1].default_value = 0.0
# Roughness
principled_bsdf.inputs[2].default_value = 0.5
# IOR
principled_bsdf.inputs[3].default_value = 1.5
# Normal
principled_bsdf.inputs[5].default_value = (0.0, 0.0, 0.0)
# Diffuse Roughness
principled_bsdf.inputs[7].default_value = 0.0
# Subsurface Weight
principled_bsdf.inputs[8].default_value = 0.0
# Subsurface Radius
principled_bsdf.inputs[9].default_value = (1.0, 0.2, 0.1)
# Subsurface Scale
principled_bsdf.inputs[10].default_value = 0.05
# Subsurface Anisotropy
principled_bsdf.inputs[12].default_value = 0.0
# Specular IOR Level
principled_bsdf.inputs[13].default_value = 0.5
# Specular Tint
principled_bsdf.inputs[14].default_value = (1.0, 1.0, 1.0, 1.0)
# Anisotropic
principled_bsdf.inputs[15].default_value = 0.0
# Anisotropic Rotation
principled_bsdf.inputs[16].default_value = 0.0
# Tangent
principled_bsdf.inputs[17].default_value = (0.0, 0.0, 0.0)
# Transmission Weight
principled_bsdf.inputs[18].default_value = 0.0
# Coat Weight
principled_bsdf.inputs[19].default_value = 0.0
# Coat Roughness
principled_bsdf.inputs[20].default_value = 0.03
# Coat IOR
principled_bsdf.inputs[21].default_value = 1.5
# Coat Tint
principled_bsdf.inputs[22].default_value = (1.0, 1.0, 1.0, 1.0)
# Coat Normal
principled_bsdf.inputs[23].default_value = (0.0, 0.0, 0.0)
# Sheen Weight
principled_bsdf.inputs[24].default_value = 0.0
# Sheen Roughness
principled_bsdf.inputs[25].default_value = 0.5
# Sheen Tint
principled_bsdf.inputs[26].default_value = (1.0, 1.0, 1.0, 1.0)
# Emission Color
principled_bsdf.inputs[27].default_value = (1.0, 1.0, 1.0, 1.0)
# Emission Strength
principled_bsdf.inputs[28].default_value = 0.0
# Thin Film Thickness
principled_bsdf.inputs[29].default_value = 0.0
# Thin Film IOR
principled_bsdf.inputs[30].default_value = 1.33
# node Attribute
attribute = n_material.nodes.new("ShaderNodeAttribute")
attribute.name = "Attribute"
attribute.attribute_name = "iso_negative"
attribute.attribute_type = "GEOMETRY"
# Set locations
material_output.location = (300.0, 300.0)
principled_bsdf.location = (32.0, 305.0)
attribute.location = (-150.0, 234.0)
# Set dimensions
material_output.width, material_output.height = 140.0, 100.0
principled_bsdf.width, principled_bsdf.height = 240.0, 100.0
attribute.width, attribute.height = 140.0, 100.0
# initialize n_material links
# principled_bsdf.BSDF -> material_output.Surface
n_material.links.new(principled_bsdf.outputs[0], material_output.inputs[0])
# attribute.Color -> principled_bsdf.Base Color
n_material.links.new(attribute.outputs[0], principled_bsdf.inputs[0])
# attribute.Alpha -> principled_bsdf.Alpha
n_material.links.new(attribute.outputs[3], principled_bsdf.inputs[4])
return mat
# auto-generated by 'Node To Python' extension
# https://extensions.blender.org/add-ons/node-to-python/
# initialize mda_density_viz node group
def _mda_density_viz_node_group():
mda_density_viz = bpy.data.node_groups.new(
type="GeometryNodeTree", name="mda_density_viz"
)
mda_density_viz.color_tag = "NONE"
mda_density_viz.description = ""
mda_density_viz.default_group_node_width = 140
# mda_density_viz interface
# Socket Geometry
geometry_socket = mda_density_viz.interface.new_socket(
name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry"
)
geometry_socket.attribute_domain = "POINT"
# Socket iso positive output
iso_positive_output_socket = mda_density_viz.interface.new_socket(
name="iso positive output", in_out="OUTPUT", socket_type="NodeSocketColor"
)
iso_positive_output_socket.default_value = (0.0, 0.0, 0.0, 1.0)
iso_positive_output_socket.default_attribute_name = "iso_positive"
iso_positive_output_socket.attribute_domain = "POINT"
# Socket iso negative output
iso_negative_output_socket = mda_density_viz.interface.new_socket(
name="iso negative output", in_out="OUTPUT", socket_type="NodeSocketColor"
)
iso_negative_output_socket.default_value = (0.0, 0.0, 0.0, 1.0)
iso_negative_output_socket.default_attribute_name = "iso_negative"
iso_negative_output_socket.attribute_domain = "POINT"
# Socket Geometry
geometry_socket_1 = mda_density_viz.interface.new_socket(
name="Geometry", in_out="INPUT", socket_type="NodeSocketGeometry"
)
geometry_socket_1.attribute_domain = "POINT"
# Socket isolevel
isolevel_socket = mda_density_viz.interface.new_socket(
name="isolevel", in_out="INPUT", socket_type="NodeSocketFloat"
)
isolevel_socket.default_value = 0.5
isolevel_socket.min_value = -1.5
isolevel_socket.max_value = 1.5
isolevel_socket.subtype = "NONE"
isolevel_socket.attribute_domain = "POINT"
isolevel_socket.force_non_field = True
# Socket iso positive
iso_positive_socket = mda_density_viz.interface.new_socket(
name="iso positive", in_out="INPUT", socket_type="NodeSocketColor"
)
iso_positive_socket.default_value = (0.0, 0.0, 1.0, 0.5)
iso_positive_socket.attribute_domain = "POINT"
iso_positive_socket.force_non_field = True
# Socket iso negative
iso_negative_socket = mda_density_viz.interface.new_socket(
name="iso negative", in_out="INPUT", socket_type="NodeSocketColor"
)
iso_negative_socket.default_value = (1.0, 0.0, 0.0, 0.5)
iso_negative_socket.attribute_domain = "POINT"
iso_negative_socket.force_non_field = True
# Socket slice X
slice_x_socket = mda_density_viz.interface.new_socket(
name="slice X", in_out="INPUT", socket_type="NodeSocketFloat"
)
slice_x_socket.default_value = 0.0
slice_x_socket.min_value = 0.0
slice_x_socket.max_value = 2.0
slice_x_socket.subtype = "NONE"
slice_x_socket.attribute_domain = "POINT"
slice_x_socket.force_non_field = True
# Socket contour
contour_socket = mda_density_viz.interface.new_socket(
name="contour", in_out="INPUT", socket_type="NodeSocketBool"
)
contour_socket.default_value = False
contour_socket.attribute_domain = "POINT"
contour_socket.force_non_field = True
# initialize mda_density_viz nodes
# node Volume to Mesh.001
volume_to_mesh_001 = mda_density_viz.nodes.new("GeometryNodeVolumeToMesh")
volume_to_mesh_001.name = "Volume to Mesh.001"
volume_to_mesh_001.resolution_mode = "GRID"
# Adaptivity
volume_to_mesh_001.inputs[4].default_value = 0.0
# node Join Geometry
join_geometry = mda_density_viz.nodes.new("GeometryNodeJoinGeometry")
join_geometry.name = "Join Geometry"
# node Group Output
group_output = mda_density_viz.nodes.new("NodeGroupOutput")
group_output.name = "Group Output"
group_output.is_active_output = True
# node Delete Geometry
delete_geometry = mda_density_viz.nodes.new("GeometryNodeDeleteGeometry")
delete_geometry.name = "Delete Geometry"
delete_geometry.domain = "POINT"
delete_geometry.mode = "ALL"
# node Compare
compare = mda_density_viz.nodes.new("FunctionNodeCompare")
compare.name = "Compare"
compare.data_type = "FLOAT"
compare.mode = "ELEMENT"
compare.operation = "LESS_EQUAL"
# node Position
position = mda_density_viz.nodes.new("GeometryNodeInputPosition")
position.name = "Position"
# node Separate XYZ
separate_xyz = mda_density_viz.nodes.new("ShaderNodeSeparateXYZ")
separate_xyz.name = "Separate XYZ"
# node Group Input
group_input = mda_density_viz.nodes.new("NodeGroupInput")
group_input.name = "Group Input"
# node Set Material
set_material = mda_density_viz.nodes.new("GeometryNodeSetMaterial")
set_material.name = "Set Material"
# Selection
set_material.inputs[1].default_value = True
set_material.inputs[2].default_value = bpy.data.materials["p_material"]
# node Set Material.001
set_material_001 = mda_density_viz.nodes.new("GeometryNodeSetMaterial")
set_material_001.name = "Set Material.001"
# Selection
set_material_001.inputs[1].default_value = True
set_material_001.inputs[2].default_value = bpy.data.materials["n_material"]
# node Volume to Mesh
volume_to_mesh = mda_density_viz.nodes.new("GeometryNodeVolumeToMesh")
volume_to_mesh.name = "Volume to Mesh"
volume_to_mesh.resolution_mode = "GRID"
# Adaptivity
volume_to_mesh.inputs[4].default_value = 0.0
# node Math
math = mda_density_viz.nodes.new("ShaderNodeMath")
math.name = "Math"
math.operation = "MULTIPLY"
math.use_clamp = False
# Value_001
math.inputs[1].default_value = -1.0
# node Set Shade Smooth
set_shade_smooth = mda_density_viz.nodes.new("GeometryNodeSetShadeSmooth")
set_shade_smooth.name = "Set Shade Smooth"
set_shade_smooth.domain = "FACE"
# Selection
set_shade_smooth.inputs[1].default_value = True
# Shade Smooth
set_shade_smooth.inputs[2].default_value = True
# node Mesh to Curve
mesh_to_curve = mda_density_viz.nodes.new("GeometryNodeMeshToCurve")
mesh_to_curve.name = "Mesh to Curve"
# node Join Geometry.001
join_geometry_001 = mda_density_viz.nodes.new("GeometryNodeJoinGeometry")
join_geometry_001.name = "Join Geometry.001"
# Set locations
volume_to_mesh_001.location = (29.0, -19.0)
join_geometry.location = (617.0, 98.0)
group_output.location = (1536.0, -106.0)
delete_geometry.location = (977.0, 132.0)
compare.location = (564.0, -209.0)
position.location = (229.0, -217.0)
separate_xyz.location = (390.0, -192.0)
group_input.location = (-484.0, -77.0)
set_material.location = (394.0, 154.0)
set_material_001.location = (394.0, 13.0)
volume_to_mesh.location = (29.0, 151.0)
math.location = (-141.0, -19.0)
set_shade_smooth.location = (795.0, 152.0)
mesh_to_curve.location = (1160.0, -144.0)
join_geometry_001.location = (1348.0, -6.0)
# Set dimensions
volume_to_mesh_001.width, volume_to_mesh_001.height = 170.0, 100.0
join_geometry.width, join_geometry.height = 140.0, 100.0
group_output.width, group_output.height = 140.0, 100.0
delete_geometry.width, delete_geometry.height = 140.0, 100.0
compare.width, compare.height = 140.0, 100.0
position.width, position.height = 140.0, 100.0
separate_xyz.width, separate_xyz.height = 140.0, 100.0
group_input.width, group_input.height = 140.0, 100.0
set_material.width, set_material.height = 140.0, 100.0
set_material_001.width, set_material_001.height = 140.0, 100.0
volume_to_mesh.width, volume_to_mesh.height = 170.0, 100.0
math.width, math.height = 140.0, 100.0
set_shade_smooth.width, set_shade_smooth.height = 140.0, 100.0
mesh_to_curve.width, mesh_to_curve.height = 140.0, 100.0
join_geometry_001.width, join_geometry_001.height = 140.0, 100.0
# initialize mda_density_viz links
# group_input.Geometry -> volume_to_mesh.Volume
mda_density_viz.links.new(group_input.outputs[0], volume_to_mesh.inputs[0])
# group_input.isolevel -> math.Value
mda_density_viz.links.new(group_input.outputs[1], math.inputs[0])
# group_input.isolevel -> volume_to_mesh.Threshold
mda_density_viz.links.new(group_input.outputs[1], volume_to_mesh.inputs[3])
# math.Value -> volume_to_mesh_001.Threshold
mda_density_viz.links.new(math.outputs[0], volume_to_mesh_001.inputs[3])
# group_input.Geometry -> volume_to_mesh_001.Volume
mda_density_viz.links.new(group_input.outputs[0], volume_to_mesh_001.inputs[0])
# volume_to_mesh.Mesh -> set_material.Geometry
mda_density_viz.links.new(volume_to_mesh.outputs[0], set_material.inputs[0])
# set_material.Geometry -> join_geometry.Geometry
mda_density_viz.links.new(set_material.outputs[0], join_geometry.inputs[0])
# volume_to_mesh_001.Mesh -> set_material_001.Geometry
mda_density_viz.links.new(volume_to_mesh_001.outputs[0], set_material_001.inputs[0])
# separate_xyz.X -> compare.A
mda_density_viz.links.new(separate_xyz.outputs[0], compare.inputs[0])
# position.Position -> separate_xyz.Vector
mda_density_viz.links.new(position.outputs[0], separate_xyz.inputs[0])
# compare.Result -> delete_geometry.Selection
mda_density_viz.links.new(compare.outputs[0], delete_geometry.inputs[1])
# group_input.slice X -> compare.B
mda_density_viz.links.new(group_input.outputs[4], compare.inputs[1])
# group_input.iso negative -> group_output.iso negative output
mda_density_viz.links.new(group_input.outputs[3], group_output.inputs[2])
# group_input.iso positive -> group_output.iso positive output
mda_density_viz.links.new(group_input.outputs[2], group_output.inputs[1])
# join_geometry_001.Geometry -> group_output.Geometry
mda_density_viz.links.new(join_geometry_001.outputs[0], group_output.inputs[0])
# set_shade_smooth.Geometry -> delete_geometry.Geometry
mda_density_viz.links.new(set_shade_smooth.outputs[0], delete_geometry.inputs[0])
# join_geometry.Geometry -> set_shade_smooth.Geometry
mda_density_viz.links.new(join_geometry.outputs[0], set_shade_smooth.inputs[0])
# delete_geometry.Geometry -> mesh_to_curve.Mesh
mda_density_viz.links.new(delete_geometry.outputs[0], mesh_to_curve.inputs[0])
# delete_geometry.Geometry -> join_geometry_001.Geometry
mda_density_viz.links.new(delete_geometry.outputs[0], join_geometry_001.inputs[0])
# group_input.contour -> mesh_to_curve.Selection
mda_density_viz.links.new(group_input.outputs[5], mesh_to_curve.inputs[1])
# set_material_001.Geometry -> join_geometry.Geometry
mda_density_viz.links.new(set_material_001.outputs[0], join_geometry.inputs[0])
# mesh_to_curve.Curve -> join_geometry_001.Geometry
mda_density_viz.links.new(mesh_to_curve.outputs[0], join_geometry_001.inputs[0])
return mda_density_viz
Raw
mda_pore_viz.py
import bpy
"""
import mda_pore_viz as p
p.draw_pore_center_line("hole1.sph")
p.create_pore("hole1.sph")
"""
def draw_pore_center_line(sphfile):
verts = []
for sph in _parse_sphfile(sphfile):
_, x, y, z, _ = sph
verts.append((x / 100, y / 100, z / 100))
edges = []
for i in range(len(verts) - 1):
edges.append((i, i + 1))
m = bpy.data.meshes.new("pore_center_line")
m.from_pydata(verts, edges, [])
m.validate()
m.update()
o = bpy.data.objects.new("pore_center_line", m)
bpy.context.collection.objects.link(o)
# This is pretty inefficient currently due to poor boolean modifier performance
def create_pore(sphfile=None):
_create_empty_pore_rgb()
for sph in _parse_sphfile(sphfile):
_, x, y, z, r = sph
color = "red"
if r >= 2.75:
color = "blue"
elif r >= 1.375:
color = "green"
# if color == "blue":
# continue
bpy.ops.mesh.primitive_uv_sphere_add(
radius=r / 100,
enter_editmode=False,
align="WORLD",
location=(x / 100, y / 100, z / 100),
scale=(1, 1, 1),
)
s = bpy.data.objects["Sphere"]
c = bpy.data.objects["pore"]
bpy.context.view_layer.objects.active = c
bpy.ops.object.modifier_add(type="BOOLEAN")
bpy.context.object.modifiers["Boolean"].operation = "UNION"
bpy.context.object.modifiers["Boolean"].object = s
slot = _get_material_slot(color)
bpy.context.object.active_material_index = slot
for i in range(slot):
bpy.ops.object.material_slot_move(direction="UP")
bpy.ops.object.modifier_apply(modifier="Boolean")
bpy.data.objects["Sphere"].select_set(True)
bpy.ops.object.delete()
def _parse_sphfile(sphfile):
spheres = []
with open(sphfile, "r") as file:
for line in file:
sline = line.strip()
if sline == "" or sline == "LAST-REC-END" or "S-888" in sline:
continue
id, x, y, z, r, _ = sline.split("SPH S")[1].split()
spheres.append([int(id), float(x), float(y), float(z), float(r)])
spheres.sort()
return spheres
def _create_empty_pore_rgb():
if "pore" in bpy.data.objects:
return
bpy.ops.mesh.primitive_cube_add(
size=2, enter_editmode=True, align="WORLD", location=(0, 0, 0), scale=(1, 1, 1)
)
bpy.context.object.name = "pore"
bpy.ops.mesh.delete(type="VERT")
bpy.ops.object.editmode_toggle()
if "red" not in bpy.data.materials:
red = bpy.data.materials.new("red")
red.use_nodes = True
bpy.data.materials["red"].node_tree.nodes["Principled BSDF"].inputs[
0
].default_value = (1, 0, 0, 1)
red = bpy.data.materials["red"]
bpy.context.object.data.materials.append(red)
if "green" not in bpy.data.materials:
red = bpy.data.materials.new("green")
red.use_nodes = True
bpy.data.materials["green"].node_tree.nodes["Principled BSDF"].inputs[
0
].default_value = (0, 1, 0, 1)
green = bpy.data.materials["green"]
bpy.context.object.data.materials.append(green)
if "blue" not in bpy.data.materials:
red = bpy.data.materials.new("blue")
red.use_nodes = True
bpy.data.materials["blue"].node_tree.nodes["Principled BSDF"].inputs[
0
].default_value = (0, 0, 1, 1)
blue = bpy.data.materials["blue"]
bpy.context.object.data.materials.append(blue)
def _get_material_slot(color):
for i in range(3):
if bpy.context.object.material_slots[i].name == color:
return i
return None
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment