<commit_before><commit_msg>Initial commit<commit_after>"""Scripts for loading molecular geometries into Blender.


    Main use is the "molecule" object, which can be used to draw models of molecular coordinates.
"""
import pickle, os, numpy as np

# Blender doesn't play well with pybel, so need to handle ImportError.
# try:
#     import pybel, openbabel
#     pybel_avail = True
# except ImportError:
#     pybel_avail = False

# Python outside of Blender doesn't play all that well with bpy, so need to handle ImportError.
try:
    import bpy
    bpy_avail = True
except ImportError:
    bpy_avail = False

# Dictionary of color definitions (RGB tuples). Blender RGB colors can be conveniently determined by
# using a uniform png of the desired color as a background image, then using the eyedropper.
COLORS = { 'white': (1.0, 1.0, 1.0),
           'light_gray': (0.7,0.7,0.7),
           'black': (0, 0, 0),
           'gray': (0.2, 0.2, 0.2), #(0.4,0.4,0.4),
           'red': (0.8, 0, 0),
           'royal_blue': (0.255, 0.412, 0.882), # This is web royal blue
           'green': (0.133, 0.545, 0.133),
           'purple': (0.627, 0.125, 0.941),
           'orange': (1.0, 0.647, 0),
           'blue_cb': (0, 0.168, 0.445),                # (0, 0.447, 0.698),
           'vermillion_cb': (0.665, 0.112, 0),          # (0.835, 0.369, 0),
           'bluish_green_cb': (0, 0.620, 0.451),
           'reddish_purple_cb': (0.800, 0.475, 0.655),
           'sky_blue_cb': (0.337, 0.706, 0.914),
           'orange_cb': (0.791, 0.347, 0),              # (0.902, 0.624, 0),
           'yellow_cb': (0.871, 0.776, 0.054)}          # 0.941, 0.894, 0.259)

# Dictionaries of colors for drawing elements, by atomic number. Used by several functions when the
# `color = 'by_element'` option is passed.
ELEMENT_COLORS = { 1: 'white',
                   6: 'gray',
                   7: 'royal_blue',
                   8: 'red' }

ELEMENT_COLORS_CB = { 1: 'white',
                      2: 'gray',
                      7: 'blue_cb',
                      8: 'vermillion_cb' }

# Conversion factors for 1 BU. Default is typically 1 nm. Assumes geometries are input with coords
# in angstroms.
UNIT_CONV = { 'nm': 0.1,
              'A': 1.0 }


ATOMIC_NUMBERS = { 'H': 1,
				   'C': 6,
				   'N': 7,
				   'O': 8 }

def _create_new_material(name, color):
    """Create a new material.

    Args:
        name (str): Name for the new material (e.g., 'red')
        color (tuple): RGB color for the new material (diffuse_color) (e.g., (1,0,0))
    Returns:
        The new material.
    """
    
    mat = bpy.data.materials.new(name)
    mat.diffuse_color = color
    mat.diffuse_shader = 'OREN_NAYAR' #'LAMBERT' 
    mat.diffuse_intensity = 0.8
    mat.roughness = 0.5
    mat.specular_color = (1, 1, 1)
    mat.specular_shader = 'BLINN' #'COOKTORR'
    mat.specular_intensity = 0.2 # 0.5
    mat.specular_hardness = 25 # 20
    mat.ambient = 1
    mat.use_transparent_shadows = True
    # mat.subsurface_scattering.use = True

    return mat


class atom:
    """A single atom.

    Attributes:
        at_num (int): The atomic number of the atom.
        location (numpy array): The xyz location of the atom, in angstroms.
        id_num (int): A unique identifier number, typically as generated by pybel.
    """

    # Dictionary of Van der Waals radii, by atomic number, from Wolfram Alpha.
    RADII = { 1: 1.20,
              6: 1.70,
              7: 1.55,
              8: 1.52}

    def __init__(self, atomic_number, location, id_num):
        self.at_num = atomic_number
        self.location = location # np.array
        self.id_num = id_num

    def draw(self, color='by_element', radius=None, units='nm', scale=1.0):
        """Draw the atom in Blender.

        Args:
            color (string, ='by_element'): If None, coloring is done by element. Otherwise specifies
                the color.
            radius (string, =None): If None, draws at the van der Waals radius. Otherwise specifies
                the radius in angstroms.
            units (sting, ='nm'): 1 BU = 1 nm by default. Can also be set to angstroms.
            scale (float, =1.0): Scaling factor for the atom. Useful when generating ball-and-stick 
                models.

        Returns:
            The blender object.
        """

        # The corrected location (i.e., scaled to units.)
        loc_corr = tuple(c * UNIT_CONV[units] for c in self.location)

        if not radius:
            bpy.ops.mesh.primitive_uv_sphere_add(
                location=loc_corr, size=atom.RADII[self.at_num]*UNIT_CONV[units]*scale)
        else:
            bpy.ops.mesh.primitive_uv_sphere_add(
                location=loc_corr, size=radius*UNIT_CONV[units]*scale)

        bpy.ops.object.shade_smooth()
        bpy.ops.object.modifier_add(type='SUBSURF')

        if color == 'by_element':
            atom_color = ELEMENT_COLORS[self.at_num]
        else:
            atom_color = color

        if not atom_color in bpy.data.materials:
            _create_new_material(atom_color, COLORS[atom_color])

        bpy.context.object.data.materials.append(bpy.data.materials[atom_color])

        bpy.context.object.name = "atom({})_{}".format(self.at_num, self.id_num)

        return bpy.context.object


class bond:
    """A bond between two atoms.

    Attributes:
        atom1 (atom): The first atom in the bond.
        atom2 (atom): The second atom in the bond.
    """

    def __init__(self, atom1, atom2):
        self.atom1 = atom1
        self.atom2 = atom2

    @staticmethod
    def _draw_half(location, length, rot_angle, rot_axis, element, radius=0.2, color='by_element',
                   units='nm'):
        """Draw half of a bond (static method).

        Draws half of a bond, given the location and length. Bonds are drawn in halves to facilitate
        coloring by element.

        Args:
            location (np.array): The center point of the half bond.
            length (float): The length of the half bond.
            rot_angle (float): Angle by which bond will be rotated.
            rot_axis (np.array): Axis of rotation.
            element (int): atomic number of element of the bond (for coloring).
            radius (float, =0.2): radius of the bond.
            color (string, ='by_element'): color of the bond. If 'by_element', uses element 
                coloring.
            units (string, ='nm'): 1 BU = 1 nm, by default. Can change to angstroms ('A').

        Returns:
            The new bond (Blender object).
        """

        loc_corr = tuple(c * UNIT_CONV[units] for c in location)
        len_corr = length * UNIT_CONV[units]
        radius_corr = radius * UNIT_CONV[units]

        bpy.ops.mesh.primitive_cylinder_add(radius=radius_corr, depth=len_corr, location=loc_corr, 
                                            end_fill_type='NOTHING')
        bpy.ops.transform.rotate(value=rot_angle, axis=rot_axis)
        bpy.ops.object.shade_smooth()
        bpy.ops.object.modifier_add(type='EDGE_SPLIT')

        if color == 'by_element':
            bond_color = ELEMENT_COLORS[element]
        else:
            bond_color = color

        if not bond_color in bpy.data.materials:
            _create_new_material(bond_color, COLORS[bond_color])

        bpy.context.object.data.materials.append(bpy.data.materials[bond_color])

        return bpy.context.object


    def draw(self, radius=0.2, color='by_element', units='nm'):
        """Draw the bond as two half bonds (to allow coloring).

        Args:
            radius (float, =0.2): Radius of cylinder in angstroms.
            color (string, ='by_element'): Color of the bond. If 'by_element', each half gets
                element coloring.
            units (string, ='nm'): 1 BU = 1 nm, by default. Can change to angstroms ('A').

        Returns:
            The bond (Blender object), with both halves joined.
        """

        created_objects = []

        center_loc = (self.atom1.location + self.atom2.location)/2
        bond_vector = self.atom1.location - self.atom2.location
        length = np.linalg.norm(bond_vector)

        bond_axis = bond_vector/length
        cyl_axis = np.array((0,0,1))
        rot_axis = np.cross(bond_axis, cyl_axis)
        angle = -np.arccos(np.dot(cyl_axis, bond_axis))

        start_center = (self.atom1.location + center_loc)/2
        created_objects.append(bond._draw_half(start_center, length/2, angle, rot_axis,
                               self.atom1.at_num, radius, color, units))

        end_center = (self.atom2.location + center_loc)/2
        created_objects.append(bond._draw_half(end_center, length/2, angle, rot_axis,
                               self.atom2.at_num, radius, color, units))

        for obj in bpy.context.selected_objects:
            obj.select = False

        for obj in created_objects:
            obj.select = True

        bpy.ops.object.join()
        bpy.context.object.name = "bond_{}({})_{}({})".format(
            self.atom1.id_num, self.atom1.at_num, self.atom2.id_num, self.atom2.at_num)
        
        return bpy.context.object


class molecule:
    """The molecule object.

    Attributes:
        atoms (list, = []): List of atoms (atom objects) in molecule.
        bonds (list, = []): List of bonds (bond objects) in molecule.
    """

    def __init__(self, name='molecule', atoms=[], bonds=[]):
        self.name = name
        self.atoms = atoms
        self.bonds = bonds
    
    def add_atom(self, atom):
        """Adds an atom to the molecule."""
        self.atoms.append(atom)

    def add_bond(self, a1id, a2id):
        """Adds a bond to the molecule, using atom ids."""
        if not self.search_bondids(a1id, a2id):
            self.bonds.append(bond(self.search_atomid(a1id), self.search_atomid(a2id)))

    def search_atomid(self, id_to_search):
        """Searches through atom list and returns atom object corresponding to (unique) id."""
        for atom in self.atoms:
            if atom.id_num == id_to_search:
                return atom
        return None

    def search_bondids(self, id1, id2):
        """Searches through bond list and returns bond object corresponding to (unique) ids."""
        for b in self.bonds:
            if ( (id1, id2) == (b.atom1.id_num, b.atom2.id_num) or 
                 (id2, id1) == (b.atom1.id_num, b.atom2.id_num) ):
                return b
        return None

    def draw_bonds(self, caps=True, radius=0.2, color='by_element', units='nm', join=True, 
                   with_H=True):
        """Draws the molecule's bonds.

        Args:
            caps (bool, =True): If true, each bond capped with sphere of radius at atom position.
                Make false if drawing ball-and-stick model using separate atom drawings.
            radius (float, =0.2): Radius of bonds in angstroms.
            color (string, ='by_element'): Color of the bonds. If 'by_element', each gets element
                coloring.
            units (string, ='nm'): 1 BU = 1 nm, by default. Can change to angstroms ('A').
            join (bool, =True): If true, all bonds are joined together into a single Bl object.

        Returns:
            The bonds as a single Blender object, if join=True. Otherwise, None.
        """
        
        created_objects = []

        for b in self.bonds:
            if with_H or ( b.atom1.at_num != 1 and b.atom2.at_num != 1 ):
                created_objects.append(b.draw(radius = radius, color = color, units = units))

        if caps:
            for a in self.atoms:
                if with_H or a.at_num != 1:
                    created_objects.append(a.draw(color = color, radius = radius, units = units))
        
        if join:
            # Deselect anything currently selected.
            for obj in bpy.context.selected_objects:
                obj.select = False

            # Select drawn bonds.
            for obj in created_objects:
                obj.select = True

            bpy.ops.object.join()

            bpy.context.object.name = self.name + '_bonds'
            
            return bpy.context.object

        else:
            return None


    def draw_atoms(self, color='by_element', radius=None, units='nm', scale=1.0, join=True, 
                   with_H=True):
        """Draw spheres for all atoms.

        Args: 
            color (str, ='by_element'): If 'by_element', uses colors in ELEMENT_COLORS. Otherwise,
                can specify color for whole model. Must be defined in COLORS.
            radius (float, =None): If specified, gives radius of all atoms.
            units (str, ='nm'): Units for 1 BU. Can also be A.
            join (bool, =True): If true, all atoms are joined together into a single Bl object.

        Returns:
            The atoms as a single Blender object, if join=True. Otherwise, None.
        """

        created_objects = [] # Holds links to all created objects, so that they can be joined.

        for a in self.atoms:
            if with_H or a.at_num != 1:
                created_objects.append(a.draw(color=color, radius=radius, units=units, scale=scale))

        if join:
            # Deselect all objects in scene.
            for obj in bpy.context.selected_objects:
                obj.select = False

            # Select all newly created objects.
            for obj in created_objects:
                obj.select = True

            bpy.ops.object.join()
            bpy.context.object.name = self.name + '_atoms'
            
            return bpy.context.object

        else:
            return None

    @classmethod
    def read_pdb(self, filename="/Users/hartlecs/blmol.pdb"):
    	"""Loads a pdb file into a molecule object. Only accepts atoms with cartesian coords through
    	the HETATM label and bonds through the CONECT label.

    	Args:
    		filename (string): The target file.

    	Returns:
    		The molecule object.

    	"""

    	newmol = molecule()

    	with open(filename) as pdbfile:
    		for line in pdbfile:
    			if line[0:6] == "HETATM":
    				idnum = int(line[6:11])
    				atnum = ATOMIC_NUMBERS[line[76:78].strip()]
    				coords = np.array((float(line[30:38]), float(line[38:46]), float(line[46:54])))
    				newmol.add_atom(atom(atnum, coords, idnum))

    			elif line[0:6] == "CONECT":

    				# Loads atoms as a list. First atom is bonded to the remaining atoms (up to
    				# four).
    				atoms = line[6:].split()
    				for bonded_atom in atoms[1:]:
    					# print(atoms[0], bonded_atom)
    					newmol.add_bond(int(atoms[0]), int(bonded_atom))
    				
    	return newmol


    # def load_pybel(self, filename="/Users/hartlecs/Dropbox/pybel_blender/oP6-Aida_AAA-2_b97d-tzv2d2p_OPT_DF.pdb", 
   	# 				 kind="pdb"):
    #     """Uses the pybel (and openbabel) modules to load in molecular geometries.

    #     Args:
    #         kind (string, ="pdb"): The type of file to be read. See the pybel documentation.
    #         filename (string): The file to be read.
    #     """

    #     all_molecules = [mol for mol in pybel.readfile(kind, filename)]

    #     for m in all_molecules:
    #         for a in m.atoms:
    #             self.add_atom(atom(atomic_number=a.atomicnum, 
    #                                location=np.array(tuple(n for n in a.coords)),
    #                                id_num=a.idx))
    #         for b in openbabel.OBMolBondIter(m.OBMol):
    #             self.add_bond(b.GetBeginAtom().GetIdx(), b.GetEndAtom().GetIdx())


    def write(self, filename="/Users/hartlecs/blender_molecule_output"):
        """Writes the object to a file using pickle.

        Args:
            filename (string): the target file.
        """

        with open(filename, 'wb') as output_file:
            pickle.dump(self, output_file)


    @classmethod
    def read(self, filename="/Users/hartlecs/blender_molecule_output"):
        """Reads in a molecule object from a file.

        Args:
            filename (string): The target file.

        Returns:
            The molecule object.
        """
        with open(filename, 'rb') as input_file:
            newmol = pickle.load(input_file)

        return newmol

{'fraction_non_alphanumeric': 0.11100832562442182, 'fraction_numerical': 0.021461609620721554, 'mean_word_length': 1.9907783105864993, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '839833', 'n_tokens_mistral': 5442, 'n_tokens_neox': 4980, 'n_words': 1471}