Skip to content

Commit

Permalink
more organization, setting up some default properties
Browse files Browse the repository at this point in the history
  • Loading branch information
@bsrunnels
bsrunnels committed Nov 21, 2023
1 parent c06b7f7 commit ed8c070
Show file tree
Hide file tree
Showing 5 changed files with 295 additions and 260 deletions.
264 changes: 5 additions & 259 deletions bleMD/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,148 +30,7 @@

from . bleMDProperties import *
from . bleMDDataFieldList import *



#
# KEY SUBROUTINE 1/2
# Opens Ovito and does basic communication with dump fil
#
def startOvito(hardrefresh=False):
filename = bpy.context.scene.bleMD_props.lammpsfile
interp = bpy.context.scene.bleMD_props.lammps_frame_stride
scene = bpy.context.scene
mytool = scene.bleMD_props

#
# Load the file
#
from ovito.io import import_file
pipeline = import_file(filename, sort_particles=True)


#
# Execute User's Ovito Python script if applicable
#
if "Ovito" in bpy.data.texts.keys():
exec(bpy.data.texts['Ovito'].as_string())

#
# Adjust number of frames
#
nframes = pipeline.source.num_frames
mytool.number_of_lammps_frames = nframes
bpy.context.scene.frame_end = nframes * mytool.lammps_frame_stride
mytool.valid_lammps_file = True

#
# Populate the properties list
#
data = pipeline.compute()
props = list(data.particles.keys())
if hardrefresh:
scene.datafieldlist.clear()

for prop in props:
if prop not in [i.name for i in scene.datafieldlist]:
item = scene.datafieldlist.add()
item.name = prop
if prop == "Position":
item.enable = True
item.editable = False

return pipeline

#
# KEY SUBROUTINE 2/2
# Updates the current data based on the Blender timestep
#
def loadUpdatedData(pipeline):
# Determine what the frame (or frames if interpolating)
# are that need to be pulled from
frame = bpy.data.scenes[0].frame_current
interp = bpy.context.scene.bleMD_props.lammps_frame_stride

# Determine interpolation (if any)
fac = (frame % interp)/interp
frame_lo = int(frame / interp)

print("FAC = ", fac)
print("frame_lo ", frame_lo)

# Set up the object or grab the existing object
# TODO: how do we handle multiple objects?
if not "MD_Object" in bpy.data.objects.keys():
# Object does not yet exist: create it
me = bpy.data.meshes.new("MD_Mesh")
ob = bpy.data.objects.new("MD_Object", me)
ob.show_name = True
bpy.context.collection.objects.link(ob)
else:
# Object exists: use it
ob = bpy.data.objects['MD_Object']
me = ob.data

# Update the data - storing the appropriate Ovito data
# in python data structure, but no updates yet.
attrs = {}
if fac == 0:
data = pipeline.compute(frame_lo)
coords = [list(xyz) for xyz in data.particles.positions]
for prop in bpy.context.scene.datafieldlist:
if prop.enable and prop.editable:
attrs[prop.name] = [x for x in data.particles[prop.name]]
#c_csym = [x for x in data.particles['c_csym']]
else:
frame_hi = frame_lo + 1
data_lo = pipeline.compute(frame_lo)
data_hi = pipeline.compute(frame_hi)
coords = [list((1-fac)*xyz_lo + fac*xyz_hi) for xyz_lo, xyz_hi in
zip(data_lo.particles.positions, data_hi.particles.positions)]
for prop in bpy.context.scene.datafieldlist:
if prop.enable and prop.editable:
attrs[prop.name] = [(1-fac)*x_lo + fac*x_hi for x_lo, x_hi in
zip(data_lo.particles[prop.name], data_hi.particles[prop.name])]

#c_csym = [(1-fac)*x_lo + fac*x_hi for x_lo,x_hi in zip(data_lo.particles['c_csym'], data_hi.particles['c_csym'])]

if not len(me.vertices):
# Do this if the object has not been created yet
# This line actually creates all the points
me.from_pydata(coords, [], [])
# Now, we go through the properties that were selected in the panel
# and set each of those properties as attributes
for prop in bpy.context.scene.datafieldlist:
if prop.enable and prop.editable:
attr = me.attributes.new(prop.name, 'FLOAT', 'POINT')
attr.data.foreach_set("value", attrs[prop.name])
else:
# We do this if we are just updating the positions and properties,
# not creating

# For some reason we have to do this in order to update the mesh
# vertex locations. There doesn't appear to be a handy blender
# routine to do this automatically
for i, v in enumerate(me.vertices):
new_location = v.co
new_location[0] = coords[i][0]
new_location[1] = coords[i][1]
new_location[2] = coords[i][2]
v.co = new_location

# Here we update the properties (e.g. c_csym)
for prop in bpy.context.scene.datafieldlist:
if prop.enable and prop.editable:
if not prop.name in me.attributes.keys():
attr = me.attributes.new(prop.name, 'FLOAT', 'POINT')
else:
attr = me.attributes.get(prop.name)
attr.data.foreach_set("value", attrs[prop.name])

me.update()

# Call setup function - Jackson
setup()
from . bleMDUtils import *


# ------------------------------------------------------------------------
Expand Down Expand Up @@ -333,6 +192,8 @@ def draw(self, context):
scene = context.scene
mytool = scene.bleMD_props

layout.prop(mytool, "override_defaults")

#
# FILE IO
#
Expand Down Expand Up @@ -395,9 +256,7 @@ def draw(self, context):

def installOvito():
import pip

pip.main(['install', 'ovito'])

import ovito

print('FINISHED')
Expand All @@ -418,7 +277,7 @@ def register():


installOvito()

import bpy.utils.previews
global custom_icons
custom_icons = bpy.utils.previews.new()
Expand All @@ -433,122 +292,9 @@ def unregister():
for cls in reversed(classes):
unregister_class(cls)
del bpy.types.Scene.bleMD_props
bpy.utils.previews.remove(custom_icons)


#
#
# Jackson messing around with automatically setting up geonode environment
#
#


def create_geonodes():
obj = bpy.data.objects["MD_Object"]

geo_nodes = obj.modifiers.get("build_geonode")
if geo_nodes:
return

geo_nodes = obj.modifiers.new("build_geonode", "NODES")

node_group = create_group()
geo_nodes.node_group = node_group


def create_group(name="geonode_object"):
group = bpy.data.node_groups.get(name)
# check if a group already exists
if group:
return

group = bpy.data.node_groups.new(name, 'GeometryNodeTree')
#group.inputs.new('NodeSocketGeometry', "Geometry")
#group.outputs.new('NodeSocketGeometry', "Geometry")
group.interface.new_socket('My Output',in_out='INPUT',socket_type='NodeSocketGeometry')
group.interface.new_socket('My Input',in_out='OUTPUT',socket_type='NodeSocketGeometry')

input_node = group.nodes.new('NodeGroupInput')
output_node = group.nodes.new('NodeGroupOutput')
output_node.is_active_output = True
input_node.location.x = -300
output_node.location.x = 350

mesh_to_points = group.nodes.new('GeometryNodeMeshToPoints')
set_material = group.nodes.new('GeometryNodeSetMaterial')
mesh_to_points.location.x = -75
set_material.location.x = 125

bpy.data.node_groups[name].nodes["Mesh to Points"].inputs[3].default_value = 1
bpy.data.node_groups["geonode_object"].nodes["Set Material"].inputs[2].default_value = bpy.data.materials["my_mat"]

group.links.new(input_node.outputs[0], mesh_to_points.inputs[0])
group.links.new(mesh_to_points.outputs[0], set_material.inputs[0])
group.links.new(set_material.outputs[0], output_node.inputs[0])

return group


def create_material():
mat = bpy.data.materials.get("my_mat")
if mat:
return

mat = bpy.data.materials.new("my_mat")
obj = bpy.data.objects["MD_Object"]
obj.data.materials.append(mat)

mat.use_nodes = True
mat_nodes = mat.node_tree.nodes
material_output = mat.node_tree.nodes.get('Material Output')
default_BSDF = mat.node_tree.nodes.get('Principled BSDF')
mat.node_tree.nodes.remove(default_BSDF)

principled = mat.node_tree.nodes.new('ShaderNodeBsdfPrincipled')
principled.location.y = 350

attribute = mat_nodes.new('ShaderNodeAttribute')
attribute.location = (-900, 250)

Math1 = mat_nodes.new('ShaderNodeMath',)
Math1.location = (-700, 250)
bpy.data.materials["my_mat"].node_tree.nodes["Math"].operation = 'SUBTRACT'
Math2 = mat_nodes.new('ShaderNodeMath',)
Math2.location = (-550, 250)
bpy.data.materials["my_mat"].node_tree.nodes["Math.001"].operation = 'DIVIDE'

color_ramp = mat_nodes.new('ShaderNodeValToRGB')
color_ramp.location = (-350, 250)

mat.node_tree.links.new(attribute.outputs[2], Math1.inputs[0])
mat.node_tree.links.new(Math1.outputs[0], Math2.inputs[0])
mat.node_tree.links.new(Math2.outputs[0], color_ramp.inputs[0])
mat.node_tree.links.new(color_ramp.outputs[0], principled.inputs[0])
mat.node_tree.links.new(principled.outputs[0], material_output.inputs[0])


def updateDefaultShader():
my_shader = bpy.context.scene.bleMD_props.my_shader
my_normalhigh = bpy.context.scene.bleMD_props.my_normalhigh
my_normallow = bpy.context.scene.bleMD_props.my_normallow
my_range = my_normalhigh - my_normallow

bpy.data.materials["my_mat"].node_tree.nodes["Attribute"].attribute_name = my_shader

bpy.data.materials["my_mat"].node_tree.nodes["Math"].inputs[1].default_value = my_normallow
bpy.data.materials["my_mat"].node_tree.nodes["Math.001"].inputs[1].default_value = my_range


return my_shader


def setup():
create_material()
create_geonodes()
for scene in bpy.data.scenes:
scene.render.engine = 'CYCLES'

print("Hello")


if __name__ == "__main__":
import bpy
Expand Down
Loading

0 comments on commit ed8c070

Please sign in to comment.