def execute(self, context):
first_atom = context.scene.objects.get(self.first_atom_name)
if first_atom and first_atom.mb.type != 'ATOM':
self.first_atom_name = " "
first_atom = None
if self.first_atom_name.strip() and not first_atom:
logger.debug('Object "{}" not found.'.format(self.first_atom_name))
return {'CANCELLED'}
if first_atom:
molecule = first_atom.mb.get_molecule()
else:
molecule = context.scene.mb.new_molecule()
# create a new atom object with the molecule's properties
new_atom = mb_utils.add_atom(context, self.coord_3d, self.element,
self.element, molecule.atom_index+1,
molecule)
self.new_atom_name = new_atom.name
# add a bond if atom is added to existing molecule
if first_atom:
new_bond = mb_utils.add_bond(context, first_atom, new_atom)
self.new_bond_name = new_bond.name
new_bond.select = True
context.scene.objects.active = new_atom
new_atom.select = True
return {'FINISHED'}
def import_molecule(context,
report,
filepath,
molecule,
refine_atoms,
refine_bonds,
bond_material,
bond_type,
auto_unit,
scale_distances,
bond_guess,
put_origin,
parent_center,
mask_planes,
mask_flip,
draw_uc,
supercell,
):
try:
all_obs = []
all_obs.append(molecule.objects.parent)
structure = mb_io_files.MB_Structure.from_file(
filepath,
auto_unit=auto_unit,
unit_fac=scale_distances,
)
# some sanity checks
if not structure.all_atoms:
msg = "No atoms found in {}. ".format(filepath)
msg += "Please check file format and/or MolBlend code."
logger.error(msg)
report({'ERROR'}, msg)
return False
logger.debug("Found {} frames in {}".format(structure.nframes,
filepath))
molecule["unit_cells"] = structure.axes
if draw_uc and molecule["unit_cells"]:
# read unit cell and create cube
unit_cell_obs = mb_utils.draw_unit_cell(molecule, context)
all_obs.extend(unit_cell_obs)
elif draw_uc and not molecule["unit_cells"]:
msg = "No unit cell vectors read."
logger.warning(msg)
report({'WARNING'}, msg)
if sum(supercell) > 3:
structure.create_supercell(supercell)
if mask_planes:
structure.apply_mask(mask_planes, mask_flip)
if bond_guess:
structure.guess_bonds(tol=0.2)
if parent_center:
center_of_mass = structure.get_center_of_mass()
else:
center_of_mass = Vector(structure.origin)
molecule.objects.parent.location = center_of_mass
# add all atoms to scene
atom_obs = {}
warning = set()
for index, atom in sorted(structure.all_atoms.items()):
new_atom = mb_utils.add_atom(context,
atom["coords"][0]-center_of_mass,
atom["element"],
atom["name"],
index, #atom["id"],
molecule)
new_atom.mb.supercell = atom.get("supercell", (0,0,0))
all_obs.append(new_atom)
new_atom.mb.index = index
atom_obs[index] = new_atom
if structure.nframes > 1:
anim_data = new_atom.animation_data_create()
atom_id = '{}.{}'.format(new_atom.mb.get_molecule().name,
new_atom.mb.index)
action = bpy.data.actions.new(name="frames_{}".format(atom_id))
anim_data.action = action
ag = action.groups.new("Location")
for dim in range(3):
fcu = action.fcurves.new(data_path="location", index=dim)
fcu.group = ag
for nf in range(structure.nframes):
coord = structure.all_atoms[index]["coords"][nf]
loc = (coord-center_of_mass)[dim]
fcu.keyframe_points.add(1)
fcu.keyframe_points[-1].co = nf + 1, loc
fcu.keyframe_points[-1].interpolation = 'LINEAR'
molecule.atom_index = index
# add bonds to scene
for index1, other in structure.bonds.items():
for index2 in other:
try:
new_bond = mb_utils.add_bond(context, atom_obs[index1],
atom_obs[index2],
bond_type=bond_type)
except KeyError as err:
if not err.args:
err.args=('',)
msg = "Atom index {} was found in CONECT"
msg += " but not in the list of atoms"
msg = msg.format(err.args[0])
err.args = (msg, )
report({'ERROR'}, msg)
raise
all_obs.append(new_bond)
molecule.bond_material = bond_material
if warning:
logger.warning('\n'.join(warning))
if put_origin:
molecule.objects.parent.location -= center_of_mass
# select all objects and make parent active
bpy.ops.object.select_all(action="DESELECT")
context.scene.objects.active = molecule.objects.parent
#for ob in all_obs:
#ob.select = True
return True
except:
# if something bad happend, delete all objects and re-raise
logger.debug("Trying to delete all newly imported objects.")
report({'ERROR'}, "Something went wrong in import. Check console.")
for ob in all_obs:
context.scene.objects.unlink(ob)
bpy.data.objects.remove(ob)
logger.exception('')
return False
def modal(self, context, event):
mouse_2d = event.mouse_x, event.mouse_y
self.coord_3d = mb_utils.mouse_2d_to_location_3d(
context, mouse_2d, region=self.region,
rv3d=self.rv3d, depth=self.depth_location)
if event.type == 'MOUSEMOVE':
new_atom = context.scene.objects.get(self.new_atom_name)
context.scene.objects.active = new_atom
new_atom.select = True
new_bond = context.scene.objects.get(self.new_bond_name)
first_atom = context.scene.objects.get(self.first_atom_name)
hover_ob = mb_utils.return_cursor_object(
context, event, exclude=[new_atom], mb_type='ATOM')
if hover_ob:
new_atom.draw_bounds_type = 'SPHERE'
new_atom.draw_type = 'BOUNDS'
if new_bond:
new_bond.constraints["mb.stretch"].target = hover_ob
else:
new_atom.draw_type = 'SOLID'
if new_bond:
new_bond.constraints["mb.stretch"].target = new_atom
if not event.alt:
self.coord_3d = mb_geometry.get_fixed_geometry(
context, first_atom, new_atom, self.coord_3d,
self.geometry)
if event.ctrl:
# constrain length
self.coord_3d = mb_geometry.get_fixed_length(
context, first_atom, new_atom, self.coord_3d,
length=-1)
new_atom.location = self.coord_3d - self.parent_coord_3d
# sometimes, when bond is exactly along axis, the dimension goes
# to zero due to the stretch constraint
# check for this case and fix it
if new_bond:
mb_utils.check_ob_dimensions(new_bond)
elif event.type == 'LEFTMOUSE' and event.value == 'RELEASE':
# check if new atom is above already existing atom
new_atom = context.object
first_atom = context.scene.objects.get(self.first_atom_name)
hover_ob = mb_utils.return_cursor_object(
context, event, exclude=[new_atom], mb_type='ATOM')
if hover_ob:
mol = new_atom.mb.get_molecule()
mol.remove_object(new_atom)
mol.atom_index -= 1
context.scene.objects.unlink(new_atom)
bpy.data.objects.remove(new_atom)
new_bond = context.scene.objects.get(self.new_bond_name)
if new_bond:
mol.remove_object(new_bond)
context.scene.objects.unlink(new_bond)
bpy.data.objects.remove(new_bond)
mb_utils.add_bond(context, first_atom, hover_ob)
return {'FINISHED'}
return {'RUNNING_MODAL'}
def import_molecule(
context,
report,
filepath,
molecule,
refine_atoms,
refine_bonds,
bond_material,
bond_type,
auto_unit,
scale_distances,
bond_guess,
put_origin,
parent_center,
mask_planes,
mask_flip,
draw_uc,
supercell,
):
try:
all_obs = []
all_obs.append(molecule.objects.parent)
structure = mb_io_files.MB_Structure.from_file(
filepath,
auto_unit=auto_unit,
unit_fac=scale_distances,
)
# some sanity checks
if not structure.all_atoms:
msg = "No atoms found in {}. ".format(filepath)
msg += "Please check file format and/or MolBlend code."
logger.error(msg)
report({'ERROR'}, msg)
return False
logger.debug("Found {} frames in {}".format(structure.nframes,
filepath))
molecule["unit_cells"] = structure.axes
if draw_uc and molecule["unit_cells"]:
# read unit cell and create cube
unit_cell_obs = mb_utils.draw_unit_cell(molecule, context)
all_obs.extend(unit_cell_obs)
elif draw_uc and not molecule["unit_cells"]:
msg = "No unit cell vectors read."
logger.warning(msg)
report({'WARNING'}, msg)
if sum(supercell) > 3:
structure.create_supercell(supercell)
if mask_planes:
structure.apply_mask(mask_planes, mask_flip)
if bond_guess:
structure.guess_bonds(tol=0.2)
if parent_center:
center_of_mass = structure.get_center_of_mass()
else:
center_of_mass = Vector(structure.origin)
molecule.objects.parent.location = center_of_mass
# add all atoms to scene
atom_obs = {}
warning = set()
for index, atom in sorted(structure.all_atoms.items()):
new_atom = mb_utils.add_atom(
context,
atom["coords"][0] - center_of_mass,
atom["element"],
atom["name"],
index, #atom["id"],
molecule)
new_atom.mb.supercell = atom.get("supercell", (0, 0, 0))
all_obs.append(new_atom)
new_atom.mb.index = index
atom_obs[index] = new_atom
if structure.nframes > 1:
anim_data = new_atom.animation_data_create()
atom_id = '{}.{}'.format(new_atom.mb.get_molecule().name,
new_atom.mb.index)
action = bpy.data.actions.new(name="frames_{}".format(atom_id))
anim_data.action = action
ag = action.groups.new("Location")
for dim in range(3):
fcu = action.fcurves.new(data_path="location", index=dim)
fcu.group = ag
for nf in range(structure.nframes):
coord = structure.all_atoms[index]["coords"][nf]
loc = (coord - center_of_mass)[dim]
fcu.keyframe_points.add(1)
fcu.keyframe_points[-1].co = nf + 1, loc
fcu.keyframe_points[-1].interpolation = 'LINEAR'
molecule.atom_index = index
# add bonds to scene
for index1, other in structure.bonds.items():
for index2 in other:
try:
new_bond = mb_utils.add_bond(context,
atom_obs[index1],
atom_obs[index2],
bond_type=bond_type)
except KeyError as err:
if not err.args:
err.args = ('', )
msg = "Atom index {} was found in CONECT"
msg += " but not in the list of atoms"
msg = msg.format(err.args[0])
err.args = (msg, )
report({'ERROR'}, msg)
raise
all_obs.append(new_bond)
molecule.bond_material = bond_material
if warning:
logger.warning('\n'.join(warning))
if put_origin:
molecule.objects.parent.location -= center_of_mass
# select all objects and make parent active
bpy.ops.object.select_all(action="DESELECT")
context.scene.objects.active = molecule.objects.parent
#for ob in all_obs:
#ob.select = True
return True
except:
# if something bad happend, delete all objects and re-raise
logger.debug("Trying to delete all newly imported objects.")
report({'ERROR'}, "Something went wrong in import. Check console.")
for ob in all_obs:
context.scene.objects.unlink(ob)
bpy.data.objects.remove(ob)
logger.exception('')
return False