Skip to content

Commit 2f400e4

Browse files
MartinBaGarMartinBaGar
authored
Merge pull request #1 from MartinBaGar/test
Test
2 parents e49f30f + 01d64d4 commit 2f400e4

File tree

7 files changed

+330
-12
lines changed

7 files changed

+330
-12
lines changed

‎molecularnodes/addon.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# along with this program. If not, see <http://www.gnu.org/licenses/>.
1313

1414
import bpy
15-
from bpy.app.handlers import frame_change_pre, load_post, save_post
15+
from bpy.app.handlers import frame_change_pre, frame_change_post, load_post, save_post
1616
from bpy.props import PointerProperty, CollectionProperty
1717
from .handlers import update_entities
1818
from . import entities, operators, props, session, ui
@@ -62,11 +62,16 @@ def register():
6262

6363
bpy.types.Scene.MNSession = session.MNSession() # type: ignore
6464
bpy.types.Object.uuid = props.uuid_property # type: ignore
65-
bpy.types.Object.mn = PointerProperty(type=props.MolecularNodesObjectProperties) # type: ignore
66-
bpy.types.Scene.mn = PointerProperty(type=props.MolecularNodesSceneProperties) # type: ignore
65+
bpy.types.Object.mn = PointerProperty(
66+
type=props.MolecularNodesObjectProperties) # type: ignore
67+
bpy.types.Scene.mn = PointerProperty(
68+
type=props.MolecularNodesSceneProperties) # type: ignore
6769
bpy.types.Object.mn_trajectory_selections = CollectionProperty( # type: ignore
6870
type=props.TrajectorySelectionItem # type: ignore
6971
)
72+
bpy.types.Scene.interaction_visualiser = PointerProperty(
73+
type=entities.interaction.interaction.InteractionVisualiserProperties
74+
)
7075

7176

7277
def unregister():
@@ -86,5 +91,6 @@ def unregister():
8691
frame_change_pre.remove(update_entities)
8792
del bpy.types.Scene.MNSession # type: ignore
8893
del bpy.types.Scene.mn # type: ignore
94+
del bpy.types.Scene.interaction_visualiser # type: ignore
8995
del bpy.types.Object.mn # type: ignore
90-
del bpy.types.Object.mn_trajectory_selections # type: ignore
96+
del bpy.types.Object.mn_trajectory_selections # type: ignore

‎molecularnodes/entities/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from . import molecule, trajectory
1+
from . import molecule, trajectory, interaction
22
from .density import MN_OT_Import_Map
33
from .trajectory.dna import MN_OT_Import_OxDNA_Trajectory
44
from .ensemble import CellPack
@@ -19,4 +19,5 @@
1919
]
2020
+ trajectory.CLASSES
2121
+ molecule.CLASSES
22+
+ interaction.CLASSES
2223
)

‎molecularnodes/entities/entity.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ class EntityType(Enum):
1313
MOLECULE = "molecule"
1414
ENSEMBLE = "ensemble"
1515
DENSITY = "density"
16+
INTERACTION = "interaction"
1617

1718

1819
class MolecularEntity(

‎molecularnodes/entities/interaction/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
from .interaction import CLASSES as INTERACTION_CLASSES
2+
from .ui import CLASSES as UI_CLASSES
3+
4+
CLASSES = INTERACTION_CLASSES + UI_CLASSES

‎molecularnodes/entities/interaction/interaction.py‎

Lines changed: 270 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,270 @@
1+
import bpy
2+
from bpy.types import Operator, PropertyGroup
3+
from bpy.props import StringProperty, FloatProperty
4+
from mathutils import Vector
5+
import json
6+
7+
from ..entity import EntityType, MolecularEntity
8+
9+
10+
class Interaction(MolecularEntity):
11+
def __init__(self):
12+
super().__init__()
13+
self._entity_type = EntityType.INTERACTION
14+
self.frame_dict = {}
15+
self.bond_objects = {}
16+
self.trajectory_name = ""
17+
self.bond_width = 0.0025
18+
19+
@staticmethod
20+
def create_bond_material(interaction_type):
21+
material_name = f"{interaction_type}_Material"
22+
mat = bpy.data.materials.get(material_name)
23+
if mat is None:
24+
mat = bpy.data.materials.new(name=material_name)
25+
mat.use_nodes = True
26+
27+
nodes = mat.node_tree.nodes
28+
nodes.clear()
29+
30+
# Material nodes
31+
texture_coord = nodes.new(type="ShaderNodeTexCoord")
32+
color_ramp = nodes.new(type="ShaderNodeValToRGB")
33+
wave_texture = nodes.new(type="ShaderNodeTexWave")
34+
principled_bsdf = nodes.new(type="ShaderNodeBsdfPrincipled")
35+
output_node = nodes.new(type="ShaderNodeOutputMaterial")
36+
37+
# Node positioning
38+
texture_coord.location = (-400, 0)
39+
wave_texture.location = (-200, 0)
40+
color_ramp.location = (0, 0)
41+
principled_bsdf.location = (200, 0)
42+
output_node.location = (400, 0)
43+
44+
links = mat.node_tree.links
45+
links.new(texture_coord.outputs["UV"],
46+
wave_texture.inputs["Vector"])
47+
links.new(wave_texture.outputs["Fac"], color_ramp.inputs["Fac"])
48+
links.new(color_ramp.outputs["Color"],
49+
principled_bsdf.inputs["Alpha"])
50+
links.new(
51+
principled_bsdf.outputs["BSDF"], output_node.inputs["Surface"])
52+
53+
# Customize based on interaction type
54+
if interaction_type == "Hydrogen":
55+
principled_bsdf.inputs["Base Color"].default_value = (
56+
0, 1, 0, 1) # Green
57+
# Tighter wave pattern
58+
wave_texture.inputs[1].default_value = 1
59+
elif interaction_type == "Salt_Bridge":
60+
principled_bsdf.inputs["Base Color"].default_value = (
61+
1, 0.5, 0, 1) # Orange
62+
# Moderate wave pattern
63+
wave_texture.inputs[1].default_value = 1
64+
elif interaction_type == "PiStacking":
65+
principled_bsdf.inputs["Base Color"].default_value = (
66+
0, 0, 1, 1) # Blue
67+
# Very tight wave pattern
68+
wave_texture.inputs[1].default_value = 1
69+
elif interaction_type == "CationPi":
70+
principled_bsdf.inputs["Base Color"].default_value = (
71+
0.5, 0, 0.5, 1) # Purple
72+
# Slightly denser wave pattern
73+
wave_texture.inputs[1].default_value = 1
74+
75+
# Wave texture settings
76+
wave_texture.wave_type = "BANDS"
77+
wave_texture.bands_direction = "X"
78+
wave_texture.wave_profile = "SIN"
79+
80+
# Wave transparency settings
81+
color_ramp.color_ramp.elements[0].position = 0.0
82+
color_ramp.color_ramp.elements[0].color = (
83+
1, 1, 1, 1) # Opaque part
84+
color_ramp.color_ramp.interpolation = "CONSTANT"
85+
color_ramp.color_ramp.elements[1].position = 0.5
86+
color_ramp.color_ramp.elements[1].color = (
87+
0, 0, 0, 0) # Transparent part
88+
89+
principled_bsdf.inputs["Alpha"].default_value = 1
90+
91+
mat.blend_method = "CLIP"
92+
mat.shadow_method = "CLIP"
93+
94+
return mat
95+
96+
@staticmethod
97+
def calculate_centroid(trajectory_object, indices):
98+
"""Calculate the centroid of a set of vertex indices."""
99+
vertices = [
100+
trajectory_object.data.vertices[int(idx)].co for idx in indices]
101+
return sum(vertices, Vector((0, 0, 0))) / len(vertices)
102+
103+
def create_bond_objects(self, all_bonds, bond_material, interaction_type):
104+
# Creates trajectory-children collections for each interaction
105+
interactions_collection = f"{self.trajectory_name}_Interactions"
106+
mn_collection = bpy.data.collections["MolecularNodes"]
107+
108+
if interactions_collection not in mn_collection.children:
109+
bonds_collection = bpy.data.collections.new(
110+
interactions_collection)
111+
mn_collection.children.link(bonds_collection)
112+
113+
interaction_collection = mn_collection.get(interaction_type)
114+
if not interaction_collection:
115+
interaction_collection = bpy.data.collections.new(interaction_type)
116+
bpy.data.collections[interactions_collection].children.link(
117+
interaction_collection)
118+
119+
# Creates bond objects that will be given coordinates on update
120+
bond_objects = {}
121+
for couple in all_bonds:
122+
# bond_name = f"{interaction_type.lower()[:6]}_{str(couple[0])}_{str(couple[1])}"
123+
bond_name = f"{str(couple[0])}_{str(couple[1])}"
124+
curve_data = bpy.data.curves.new(name=bond_name, type="CURVE")
125+
curve_object = bpy.data.objects.new(bond_name, curve_data)
126+
interaction_collection.objects.link(curve_object)
127+
curve_data.dimensions = "3D"
128+
curve_data.bevel_depth = self.bond_width
129+
spline = curve_data.splines.new(type="POLY")
130+
spline.points.add(1)
131+
curve_object.data.materials.append(bond_material)
132+
bond_objects[couple] = curve_object
133+
134+
return bond_objects
135+
136+
def update_bond_positions(self, scene):
137+
blender_object = bpy.data.objects.get(self.trajectory_name)
138+
if not blender_object or not blender_object.data:
139+
return
140+
141+
frame = str(scene)
142+
for interaction_type, type_bonds in self.bond_objects.items():
143+
couples = self.frame_dict[interaction_type].get(frame, set())
144+
145+
for bond, obj in type_bonds.items():
146+
if obj and obj.data:
147+
visible = bond in couples
148+
obj.hide_viewport = obj.hide_render = not visible
149+
150+
if visible:
151+
spline = obj.data.splines[0]
152+
153+
if interaction_type == "PiStacking":
154+
lig_centroid = self.calculate_centroid(
155+
blender_object, bond[0])
156+
prot_centroid = self.calculate_centroid(
157+
blender_object, bond[1])
158+
spline.points[0].co = (*lig_centroid, 1)
159+
spline.points[1].co = (*prot_centroid, 1)
160+
else:
161+
spline.points[0].co = (
162+
*blender_object.data.vertices[int(bond[0])].co, 1)
163+
spline.points[1].co = (
164+
*blender_object.data.vertices[int(bond[1])].co, 1)
165+
166+
def setup_from_json(self, json_file, object_name):
167+
self.trajectory_name = object_name
168+
169+
with open(json_file, "r") as file:
170+
hb_data = json.load(file)
171+
172+
# Setup frame_dict
173+
for interaction_type, frames in hb_data.items():
174+
if interaction_type not in self.frame_dict:
175+
self.frame_dict[interaction_type] = {}
176+
177+
for frame_key, interactions in frames.items():
178+
if frame_key not in self.frame_dict[interaction_type]:
179+
self.frame_dict[interaction_type][frame_key] = set()
180+
181+
for interaction in interactions:
182+
if interaction_type == "PiStacking":
183+
self.frame_dict[interaction_type][frame_key].add(
184+
(tuple(interaction["Ligand"]),
185+
tuple(interaction["Protein"]))
186+
)
187+
else:
188+
self.frame_dict[interaction_type][frame_key].add(
189+
(float(interaction["Ligand"]),
190+
float(interaction["Protein"]))
191+
)
192+
193+
# Create bond objects
194+
for interaction_type, frames in self.frame_dict.items():
195+
all_bonds = set().union(
196+
*self.frame_dict[interaction_type].values())
197+
bond_material = self.create_bond_material(interaction_type)
198+
self.bond_objects[interaction_type] = self.create_bond_objects(
199+
all_bonds, bond_material, interaction_type
200+
)
201+
202+
def update_bevel_depth(self, new_depth):
203+
self.bond_width = new_depth
204+
for interaction_type, bond_objects in self.bond_objects.items():
205+
for bond, obj in bond_objects.items():
206+
if obj and obj.type == 'CURVE':
207+
obj.data.bevel_depth = new_depth*0.001
208+
209+
210+
def update_bevel_depth(self, context):
211+
new_depth = self.bond_width # Get the updated value
212+
scene = context.scene
213+
214+
# Loop through the entities and update their bevel depth
215+
for obj in scene.MNSession.entities.values():
216+
if obj._entity_type.value == "interaction":
217+
obj.update_bevel_depth(new_depth)
218+
219+
220+
class OBJECT_OT_interaction_visualiser(Operator):
221+
bl_idname = "object.interaction_visualiser"
222+
bl_label = "Visualize Interactions"
223+
bl_options = {"REGISTER", "UNDO"}
224+
225+
def execute(self, context):
226+
scene = context.scene
227+
interaction_props = scene.interaction_visualiser
228+
json_file = interaction_props.json_file
229+
230+
# Create and setup interaction
231+
interaction = Interaction()
232+
233+
# Store the active object's name
234+
blender_object = context.active_object
235+
if blender_object:
236+
object_name = blender_object.name
237+
interaction.setup_from_json(json_file, object_name)
238+
239+
self.report({"INFO"}, "Interaction visualization setup complete")
240+
return {"FINISHED"}
241+
else:
242+
self.report({"ERROR"}, "No active object selected")
243+
return {"CANCELLED"}
244+
245+
246+
class InteractionVisualiserProperties(PropertyGroup):
247+
json_file: StringProperty(
248+
name="JSON File",
249+
description="Path to the JSON file containing data",
250+
default="",
251+
maxlen=1024,
252+
subtype="FILE_PATH",
253+
)
254+
object_name: StringProperty(
255+
name="Blender Object Name",
256+
description="Name of the Blender object representing the molecule",
257+
default="",
258+
)
259+
bond_width: FloatProperty(
260+
name="Bond Width",
261+
description="Update the bevel depth of the interaction objects",
262+
default=2.5,
263+
update=update_bevel_depth,
264+
)
265+
266+
267+
CLASSES = [
268+
InteractionVisualiserProperties,
269+
OBJECT_OT_interaction_visualiser,
270+
]

‎molecularnodes/entities/interaction/ui.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
from bpy.types import Panel
2+
3+
4+
class VIEW3D_PT_Interactions(Panel):
5+
bl_label = "MN Interactions"
6+
bl_space_type = "VIEW_3D"
7+
bl_region_type = "UI"
8+
bl_category = "MN Interactions"
9+
10+
def draw(self, context):
11+
layout = self.layout
12+
scene = context.scene
13+
interaction_props = scene.interaction_visualiser
14+
15+
layout.prop(interaction_props, "json_file")
16+
layout.prop(interaction_props, "object_name")
17+
layout.operator("object.interaction_visualiser")
18+
19+
20+
class VIEW3D_PT_Interactions_Customisation(Panel):
21+
bl_label = "Customise bonds"
22+
bl_idname = "IV_subpanel_1"
23+
bl_space_type = 'VIEW_3D'
24+
bl_region_type = 'UI'
25+
bl_category = "Custom Tab"
26+
bl_parent_id = "VIEW3D_PT_Interactions"
27+
28+
def draw(self, context):
29+
layout = self.layout
30+
scene = context.scene
31+
interaction_props = scene.interaction_visualiser
32+
33+
layout.label(text="Characteristics")
34+
layout.prop(interaction_props, "bond_width")
35+
layout.label(text="Material")
36+
# layout.prop(interaction_props, "material")
37+
# layout.prop(interaction_props, "material")
38+
39+
40+
CLASSES = [
41+
VIEW3D_PT_Interactions,
42+
VIEW3D_PT_Interactions_Customisation,
43+
]

‎molecularnodes/handlers.py‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,3 @@ def update_entities(scene):
5050
entity.set_frame(frame_to_set)
5151
except NotImplementedError:
5252
pass
53-
54-
else:
55-
# this is the old method of updating the trajectories and is maintained for
56-
# backwards compatibility # TODO: takeout for later release
57-
entity._update_positions(scene.frame_current)
58-
entity._update_selections()
59-
entity._update_calculations()

0 commit comments

Comments
 (0)