pdb_import.py

                FLAG_s2 = True
                break
            s2 += 1

        # If the vertex (atom) is not yet in the vertex list:
        # append the number of atom and the vertex to the two lists.
        # For the first atom:
        if FLAG_s1 == False:
            atom1 = copy(all_atoms[stick.atom1-1].location)
            stick_vertices.append(atom1)
            stick_vertices_nr.append(stick.atom1-1)
        # For the second atom:
        if FLAG_s2 == False:
            atom2 = copy(all_atoms[stick.atom2-1].location)
            stick_vertices.append(atom2)
            stick_vertices_nr.append(stick.atom2-1)

        # Build the edges:

        # If both vertices (atoms) were not in the lists, then
        # the edge is simply [i,i+1]. These are two new vertices
        # (atoms), so increase i by 2.
        if FLAG_s1 == False and FLAG_s2 == False:
            stick_edges.append([i,i+1])
            i += 2
        # Both vertices (atoms) were already in the list, so then
        # use the vertices (atoms), which already exist. They are
        # at positions s1 and s2.
        if FLAG_s1 == True and FLAG_s2 == True:
            stick_edges.append([s1,s2])
        # The following two if cases describe the situation that
        # only one vertex (atom) was in the list. Since only ONE
        # new vertex was added, increase i by one.
        if FLAG_s1 == True and FLAG_s2 == False:
            stick_edges.append([s1,i])
            i += 1
        if FLAG_s1 == False and FLAG_s2 == True:
            stick_edges.append([i,s2])
            i += 1

    # Build the mesh of the sticks
    stick_mesh = bpy.data.meshes.new("Mesh_sticks")
    stick_mesh.from_pydata(stick_vertices, stick_edges, [])
    stick_mesh.update()
    new_stick_mesh = bpy.data.objects.new("Sticks", stick_mesh)
    # Link the active mesh to the molecule collection
    coll_molecule.objects.link(new_stick_mesh)

    # Apply the skin modifier.
    new_stick_mesh.modifiers.new(name="Sticks_skin", type='SKIN')
    # Smooth the skin surface if this option has been chosen.
    new_stick_mesh.modifiers[0].use_smooth_shade = use_sticks_smooth
    # Apply the Subdivision modifier.
    new_stick_mesh.modifiers.new(name="Sticks_subsurf", type='SUBSURF')
    # Options: choose the levels
    new_stick_mesh.modifiers[1].levels = sticks_subdiv_view
    new_stick_mesh.modifiers[1].render_levels = sticks_subdiv_render

    stick_material = bpy.data.materials.new(ELEMENTS[-1].name)
    stick_material.diffuse_color = ELEMENTS[-1].color
    new_stick_mesh.active_material = stick_material

    # This is for putting the radius of the sticks onto
    # the desired value 'Stick_diameter'
    bpy.context.view_layer.objects.active = new_stick_mesh
    # EDIT mode
    bpy.ops.object.mode_set(mode='EDIT', toggle=False)
    bm = bmesh.from_edit_mesh(new_stick_mesh.data)
    bpy.ops.mesh.select_all(action='DESELECT')

    # Select all vertices
    for v in bm.verts:
        v.select = True

    # This is somewhat a factor for the radius.
    r_f = 4.0
    # Apply operator 'skin_resize'.
    bpy.ops.transform.skin_resize(
        value=(
            Stick_diameter * r_f,
            Stick_diameter * r_f,
            Stick_diameter * r_f,
        ),
        constraint_axis=(False, False, False),
        orient_type='GLOBAL',
        mirror=False,
        use_proportional_edit=False,
        snap=False,
        snap_target='CLOSEST',
        snap_point=(0, 0, 0),
        snap_align=False,
        snap_normal=(0, 0, 0),
        release_confirm=False,
    )
    # Back to the OBJECT mode.
    bpy.ops.object.mode_set(mode='OBJECT', toggle=False)

    return new_stick_mesh


# Draw the sticks the normal way: connect the atoms by simple cylinders.
# Two options: 1. single cylinders parented to an empty
#              2. one single mesh object
def draw_sticks_normal(all_atoms,
                       all_sticks,
                       center,
                       Stick_diameter,
                       Stick_sectors,
                       use_sticks_smooth,
                       use_sticks_one_object,
                       use_sticks_one_object_nr,
                       coll_molecule):

    stick_material = bpy.data.materials.new(ELEMENTS[-1].name)
    stick_material.diffuse_color = ELEMENTS[-1].color

    up_axis = Vector([0.0, 0.0, 1.0])

    # For all sticks, do ...
    list_group = []
    list_group_sub = []
    counter = 0
    for stick in all_sticks:

        # The vectors of the two atoms
        atom1 = all_atoms[stick.atom1-1].location-center
        atom2 = all_atoms[stick.atom2-1].location-center
        # Location
        location = (atom1 + atom2) * 0.5
        # The difference of both vectors
        v = (atom2 - atom1)
        # Angle with respect to the z-axis
        angle = v.angle(up_axis, 0)
        # Cross-product between v and the z-axis vector. It is the
        # vector of rotation.
        axis = up_axis.cross(v)
        # Calculate Euler angles
        euler = Matrix.Rotation(angle, 4, axis).to_euler()
        # Create stick
        stick = bpy.ops.mesh.primitive_cylinder_add(vertices=Stick_sectors,
                                                    radius=Stick_diameter,
                                                    depth=v.length,
                                                    end_fill_type='NGON',
                                                    align='WORLD',
                                                    enter_editmode=False,
                                                    location=location,
                                                    rotation=(0, 0, 0))
        # Put the stick into the scene ...
        stick = bpy.context.view_layer.objects.active
        # ... and rotate the stick.
        stick.rotation_euler = euler
        # ... and name
        stick.name = "Stick_Cylinder"
        counter += 1

        # Smooth the cylinder.
        if use_sticks_smooth == True:
            bpy.ops.object.select_all(action='DESELECT')
            stick.select_set(True)
            bpy.ops.object.shade_smooth()

        list_group_sub.append(stick)

        if use_sticks_one_object == True:
            if counter == use_sticks_one_object_nr:
                bpy.ops.object.select_all(action='DESELECT')
                for stick in list_group_sub:
                    stick.select_set(True)
                bpy.ops.object.join()
                list_group.append(bpy.context.view_layer.objects.active)
                bpy.ops.object.select_all(action='DESELECT')
                list_group_sub = []
                counter = 0
        else:
            # Material ...
            stick.active_material = stick_material

    if use_sticks_one_object == True:
        bpy.ops.object.select_all(action='DESELECT')
        for stick in list_group_sub:
            stick.select_set(True)
        bpy.ops.object.join()
        list_group.append(bpy.context.view_layer.objects.active)
        bpy.ops.object.select_all(action='DESELECT')

        for group in list_group:
            group.select_set(True)
        bpy.ops.object.join()
        bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY',
                                   center='MEDIAN')
        sticks = bpy.context.view_layer.objects.active
        sticks.active_material = stick_material

        sticks.location += center

        # Collections
        # ===========
        # Note the collection where the sticks were placed into.
        coll_all = sticks.users_collection
        if len(coll_all) > 0:
            coll_past = coll_all[0]
        else:
            coll_past = bpy.context.scene.collection

        # Link the sticks with the collection of the molecule ...
        coll_molecule.objects.link(sticks)
        # ... and unlink them from the collection it has been before.
        coll_past.objects.unlink(sticks)

        return sticks
    else:
        # Here we use an empty ...
        bpy.ops.object.empty_add(type='ARROWS',
                                  align='WORLD',
                                  location=(0, 0, 0),
                                  rotation=(0, 0, 0))
        sticks_empty = bpy.context.view_layer.objects.active
        sticks_empty.name = "A_sticks_empty"
        # ... that is parent to all sticks. With this, we can better move
        # all sticks if necessary.
        for stick in list_group_sub:
            stick.parent = sticks_empty

        sticks_empty.location += center

        # Collections
        # ===========
        # Create a collection that will contain all sticks + the empty and ...
        coll = bpy.data.collections.new("Sticks")
        # ... link it to the collection, which contains all parts of the
        # molecule.
        coll_molecule.children.link(coll)
        # Now, create a collection that only contains the sticks and ...
        coll_cylinder = bpy.data.collections.new("Sticks_cylinders")
        # ... link it to the collection, which contains the sticks and empty.
        coll.children.link(coll_cylinder)

        # Note the collection where the empty was placed into, ...
        coll_all = sticks_empty.users_collection
        if len(coll_all) > 0:
            coll_past = coll_all[0]
        else:
            coll_past = bpy.context.scene.collection
        # ... link the empty with the new collection  ...
        coll.objects.link(sticks_empty)
        # ... and unlink it from the old collection where it has been before.
        coll_past.objects.unlink(sticks_empty)

        # Note the collection where the cylinders were placed into, ...
        coll_all = list_group_sub[0].users_collection
        if len(coll_all) > 0:
            coll_past = coll_all[0]
        else:
            coll_past = bpy.context.scene.collection

        for stick in list_group_sub:
            # ... link each stick with the new collection  ...
            coll_cylinder.objects.link(stick)
            # ... and unlink it from the old collection.
            coll_past.objects.unlink(stick)

        return sticks_empty


# -----------------------------------------------------------------------------
#                                                            The main routine

def import_pdb(Ball_type,
               Ball_azimuth,
               Ball_zenith,
               Ball_radius_factor,
               radiustype,
               Ball_distance_factor,
               use_sticks,
               use_sticks_type,
               sticks_subdiv_view,
               sticks_subdiv_render,
               use_sticks_color,
               use_sticks_smooth,
               use_sticks_bonds,
               use_sticks_one_object,
               use_sticks_one_object_nr,
               Stick_unit, Stick_dist,
               Stick_sectors,
               Stick_diameter,
               put_to_center,
               use_camera,
               use_light,
               filepath_pdb):

    # List of materials
    atom_material_list = []

    # A list of ALL objects which are loaded (needed for selecting the loaded
    # structure.
    atom_object_list = []

    # ------------------------------------------------------------------------
    # INITIALIZE THE ELEMENT LIST

    read_elements()

    # ------------------------------------------------------------------------
    # READING DATA OF ATOMS

    (Number_of_total_atoms, all_atoms) = read_pdb_file(filepath_pdb, radiustype)

    # ------------------------------------------------------------------------
    # MATERIAL PROPERTIES FOR ATOMS

    # The list that contains info about all types of atoms is created
    # here. It is used for building the material properties for
    # instance (see below).
    atom_all_types_list = []

    for atom in all_atoms:
        FLAG_FOUND = False
        for atom_type in atom_all_types_list:
            # If the atom name is already in the list, FLAG on 'True'.
            if atom_type[0] == atom.name:
                FLAG_FOUND = True
                break
        # No name in the current list has been found? => New entry.
        if FLAG_FOUND == False:
            # Stored are: Atom label (e.g. 'Na'), the corresponding atom
            # name (e.g. 'Sodium') and its color.
            atom_all_types_list.append([atom.name, atom.element, atom.color])

    # The list of materials is built.
    # Note that all atoms of one type (e.g. all hydrogens) get only ONE
    # material! This is good because then, by activating one atom in the
    # Blender scene and changing the color of this atom, one changes the color
    # of ALL atoms of the same type at the same time.

    # Create first a new list of materials for each type of atom
    # (e.g. hydrogen)
    for atom_type in atom_all_types_list:
        material = bpy.data.materials.new(atom_type[1])
        material.name = atom_type[0]
        material.diffuse_color = atom_type[2]
        atom_material_list.append(material)

    # Now, we go through all atoms and give them a material. For all atoms ...
    for atom in all_atoms:
        # ... and all materials ...
        for material in atom_material_list:
            # ... select the correct material for the current atom via
            # comparison of names ...
            if atom.name in material.name:
                # ... and give the atom its material properties.
                # However, before we check, if it is a vacancy, because then it
                # gets some additional preparation. The vacancy is represented
                # by a transparent cube.
                if atom.name == "Vacancy":
                    # Some properties for eevee.
                    material.metallic = 0.8
                    material.specular_intensity = 0.5
                    material.roughness = 0.3
                    material.blend_method = 'OPAQUE'
                    material.show_transparent_back = False
                    # Some properties for cycles
                    material.use_nodes = True
                    mat_P_BSDF = material.node_tree.nodes['Principled BSDF']
                    mat_P_BSDF.inputs['Metallic'].default_value = 0.1
                    mat_P_BSDF.inputs['Roughness'].default_value = 0.2
                    mat_P_BSDF.inputs['Transmission'].default_value = 0.97
                    mat_P_BSDF.inputs['IOR'].default_value = 0.8
                # The atom gets its properties.
                atom.material = material

    # ------------------------------------------------------------------------
    # READING DATA OF STICKS

    all_sticks = read_pdb_file_sticks(filepath_pdb,
                                      use_sticks_bonds,
                                      all_atoms)
    #
    # So far, all atoms, sticks and materials have been registered.
    #

    # ------------------------------------------------------------------------
    # TRANSLATION OF THE STRUCTURE TO THE ORIGIN

    # It may happen that the structure in a PDB file already has an offset
    # If chosen, the structure is first put into the center of the scene
    # (the offset is subtracted).

    if put_to_center == True:
        sum_vec = Vector((0.0,0.0,0.0))
        # Sum of all atom coordinates
        sum_vec = sum([atom.location for atom in all_atoms], sum_vec)
        # Then the average is taken
        sum_vec = sum_vec / Number_of_total_atoms
        # After, for each atom the center of gravity is subtracted
        for atom in all_atoms:
            atom.location -= sum_vec

    # ------------------------------------------------------------------------
    # SCALING

    # Take all atoms and adjust their radii and scale the distances.
    for atom in all_atoms:
        atom.location *= Ball_distance_factor

    # ------------------------------------------------------------------------
    # DETERMINATION OF SOME GEOMETRIC PROPERTIES

    # In the following, some geometric properties of the whole object are
    # determined: center, size, etc.
    sum_vec = Vector((0.0,0.0,0.0))

    # First the center is determined. All coordinates are summed up ...
    sum_vec = sum([atom.location for atom in all_atoms], sum_vec)

    # ... and the average is taken. This gives the center of the object.
    object_center_vec = sum_vec / Number_of_total_atoms

    # Now, we determine the size.The farthest atom from the object center is
    # taken as a measure. The size is used to place well the camera and light
    # into the scene.
    object_size_vec = [atom.location - object_center_vec for atom in all_atoms]
    object_size = max(object_size_vec).length

    # ------------------------------------------------------------------------
    # SORTING THE ATOMS

    # Lists of atoms of one type are created. Example:
    # draw_all_atoms = [ data_hydrogen,data_carbon,data_nitrogen ]
    # data_hydrogen = [["Hydrogen", Material_Hydrogen, Vector((x,y,z)), 109], ...]

    # Go through the list which contains all types of atoms. It is the list,
    # which has been created on the top during reading the PDB file.
    # Example: atom_all_types_list = ["hydrogen", "carbon", ...]
    draw_all_atoms = []
    for atom_type in atom_all_types_list:

        # Don't draw 'TER atoms'.
        if atom_type[0] == "TER":
            continue

        # This is the draw list, which contains all atoms of one type (e.g.
        # all hydrogens) ...
        draw_all_atoms_type = []

        # Go through all atoms ...
        for atom in all_atoms:
            # ... select the atoms of the considered type via comparison ...
            if atom.name == atom_type[0]:
                # ... and append them to the list 'draw_all_atoms_type'.
                draw_all_atoms_type.append([atom.name,
                                            atom.material,
                                            atom.location,
                                            atom.radius])

        # Now append the atom list to the list of all types of atoms
        draw_all_atoms.append(draw_all_atoms_type)

    # ------------------------------------------------------------------------
    # COLLECTION

    # Before we start to draw the atoms and sticks, we first create a
    # collection for the molecule. All atoms (balls) and sticks (cylinders)
    # are put into this collection.
    coll_molecule_name = os.path.basename(filepath_pdb)
    scene = bpy.context.scene
    coll_molecule = bpy.data.collections.new(coll_molecule_name)
    scene.collection.children.link(coll_molecule)

    # ------------------------------------------------------------------------
    # DRAWING THE ATOMS

    bpy.ops.object.select_all(action='DESELECT')

    list_coll_elements = []
    # For each list of atoms of ONE type (e.g. Hydrogen)
    for draw_all_atoms_type in draw_all_atoms:

        atom_mesh, coll_element = draw_atoms_one_type(draw_all_atoms_type,
                                                      Ball_type,
                                                      Ball_azimuth,
                                                      Ball_zenith,
                                                      Ball_radius_factor,
                                                      object_center_vec,
                                                      coll_molecule)
        atom_object_list.append(atom_mesh)
        list_coll_elements.append(coll_element)

    # ------------------------------------------------------------------------
    # DRAWING THE STICKS: cylinders in a dupliverts structure

    if use_sticks == True and use_sticks_type == '0' and all_sticks != []:

        sticks = draw_sticks_dupliverts(all_atoms,
                                        atom_all_types_list,
                                        object_center_vec,
                                        all_sticks,
                                        Stick_diameter,
                                        Stick_sectors,
                                        Stick_unit,
                                        Stick_dist,
                                        use_sticks_smooth,
                                        use_sticks_color,
                                        list_coll_elements)
        for stick in sticks:
            atom_object_list.append(stick)

    # ------------------------------------------------------------------------
    # DRAWING THE STICKS: skin and subdivision modifier

    if use_sticks == True and use_sticks_type == '1' and all_sticks != []:

        sticks = draw_sticks_skin(all_atoms,
                                  all_sticks,
                                  Stick_diameter,
                                  use_sticks_smooth,
                                  sticks_subdiv_view,
                                  sticks_subdiv_render,
                                  coll_molecule)
        atom_object_list.append(sticks)

    # ------------------------------------------------------------------------
    # DRAWING THE STICKS: normal cylinders

    if use_sticks == True and use_sticks_type == '2' and all_sticks != []:

        sticks = draw_sticks_normal(all_atoms,
                                    all_sticks,
                                    object_center_vec,
                                    Stick_diameter,
                                    Stick_sectors,
                                    use_sticks_smooth,
                                    use_sticks_one_object,
                                    use_sticks_one_object_nr,
                                    coll_molecule)
        atom_object_list.append(sticks)

    # ------------------------------------------------------------------------
    # CAMERA and LIGHT SOURCES

    camera_light_source(use_camera,
                        use_light,
                        object_center_vec,
                        object_size)

    # ------------------------------------------------------------------------
    # SELECT ALL LOADED OBJECTS
    bpy.ops.object.select_all(action='DESELECT')
    obj = None
    for obj in atom_object_list:
        obj.select_set(True)

    # activate the last selected object
    if obj:
        bpy.context.view_layer.objects.active = obj