vedo.chemistry

  1import numpy as np
  2from vtkmodules.util.numpy_support import vtk_to_numpy
  3from vtkmodules.vtkCommonDataModel import vtkMolecule
  4from vtkmodules.vtkFiltersCore import vtkMoleculeAppend
  5from vtkmodules.vtkIOChemistry import vtkPDBReader
  6from vtkmodules.vtkRenderingCore import vtkActor, vtkPolyDataMapper
  7from vtkmodules.vtkDomainsChemistry import vtkPeriodicTable, vtkMoleculeMapper
  8from vtkmodules.vtkDomainsChemistry import vtkProteinRibbonFilter
  9from vtkmodules.vtkCommonCore import vtkIdList
 10from vtkmodules.vtkCommonDataModel import vtkPolyData
 11
 12__doc__ = """
 13This module provides a Vedo-compatible interface to the periodic table of elements,
 14allowing access to element properties such as atomic numbers, names, symbols, and covalent radiii.
 15It wraps the VTK's vtkPeriodicTable class to provide a more Pythonic interface.
 16"""
 17
 18__docformat__ = "google"
 19
 20class PeriodicTable:
 21    """
 22    A Vedo-compatible class for accessing periodic table data, wrapping vtkPeriodicTable.
 23
 24    This class provides access to element properties such as atomic numbers, names,
 25    symbols, and covalent radii, using VTK's built-in periodic table database.
 26
 27    Attributes:
 28        periodic_table (vtkPeriodicTable): The underlying VTK periodic table object.
 29    """
 30
 31    def __init__(self):
 32        """
 33        Initialize the PeriodicTable with VTK's built-in periodic table data.
 34        """
 35        self.periodic_table = vtkPeriodicTable()
 36
 37    def get_element_name(self, atomic_number):
 38        """
 39        Get the name of the element with the given atomic number.
 40
 41        Arguments:
 42            atomic_number : (int)
 43                The atomic number of the element.
 44
 45        Returns:
 46            str: The name of the element.
 47        """
 48        return self.periodic_table.GetElementName(atomic_number)
 49
 50    def get_element_symbol(self, atomic_number):
 51        """
 52        Get the symbol of the element with the given atomic number.
 53
 54        Arguments:
 55            atomic_number : (int)
 56                The atomic number of the element.
 57
 58        Returns:
 59            The symbol of the element.
 60        """
 61        return self.periodic_table.GetSymbol(atomic_number)
 62
 63    def get_atomic_number(self, symbol):
 64        """
 65        Get the atomic number of the element with the given symbol.
 66
 67        Arguments:
 68            symbol : (str)
 69                The symbol of the element.
 70
 71        Returns:
 72            The atomic number of the element.
 73        """
 74        return self.periodic_table.GetAtomicNumber(symbol)
 75
 76    def get_covalent_radius(self, atomic_number):
 77        """
 78        Get the covalent radius of the element with the given atomic number.
 79
 80        Arguments:
 81            atomic_number : (int)
 82                The atomic number of the element.
 83
 84        Returns:
 85            The covalent radius of the element.
 86        """
 87        return self.periodic_table.GetCovalentRadius(atomic_number)
 88
 89    def get_vdw_radius(self, atomic_number):
 90        """
 91        Get the Van der Waals radius of the element with the given atomic number.
 92
 93        Arguments:
 94            atomic_number: (int)
 95                The atomic number of the element.
 96
 97        Returns:
 98            The Van der Waals radius of the element.
 99        """
100        return self.periodic_table.GetVDWRadius(atomic_number)
101
102    def get_number_of_elements(self):
103        """
104        Get the total number of elements in the periodic table.
105
106        Returns:
107            The number of elements.
108        """
109        return self.periodic_table.GetNumberOfElements()
110
111    def get_element_data(self, atomic_number):
112        """
113        Get all data for the element with the given atomic number.
114
115        Arguments:
116            atomic_number : (int)
117                The atomic number of the element.
118
119        Returns:
120            A dictionary containing the element's name, symbol, and radii.
121        """
122        return {
123            "name": self.get_element_name(atomic_number),
124            "symbol": self.get_element_symbol(atomic_number),
125            "covalent_radius": self.get_covalent_radius(atomic_number),
126            "vdw_radius": self.get_vdw_radius(atomic_number),
127        }
128
129    def __getitem__(self, atomic_number):
130        """
131        Get element data by atomic number.
132
133        Arguments:
134            atomic_number: (int)
135                The atomic number of the element.
136
137        Returns:
138            A dictionary containing the element's name, symbol, and radii.
139        """
140        return self.get_element_data(atomic_number)
141
142    def __len__(self):
143        """
144        Get the number of elements in the periodic table.
145
146        Returns:
147            The number of elements.
148        """
149        return self.get_number_of_elements()
150
151    def __iter__(self):
152        """
153        Iterate over all elements in the periodic table.
154
155        Yields:
156            tuple: (atomic_number, element_data) for each element.
157        """
158        for atomic_number in range(1, self.get_number_of_elements() + 1):
159            yield atomic_number, self.get_element_data(atomic_number)
160
161    def __contains__(self, atomic_number):
162        """
163        Check if an atomic number is in the periodic table.
164
165        Arguments:
166            atomic_number: (int)
167                The atomic number to check.
168
169        Returns:
170            bool: True if the atomic number exists, False otherwise.
171        """
172        return 1 <= atomic_number <= self.get_number_of_elements()
173
174    def __str__(self):
175        """Print info about the periodic table."""
176        from vedo.colors import printc
177        module = self.__class__.__module__
178        name = self.__class__.__name__
179        out = printc(
180            f"{module}.{name} at ({hex(id(self))})".ljust(75),
181            bold=True, invert=True, return_string=True,
182        )
183        out += "\x1b[0m"
184        n = self.get_number_of_elements()
185        out += f"Number of elements  : {n}\n"
186        # Example usage of the periodic table
187        atomic_number = 12
188        out += f"Atomic number       : {atomic_number} (example entry)\n"
189        out += f"Element name        : {self.get_element_name(atomic_number)}\n"
190        out += f"Element symbol      : {self.get_element_symbol(atomic_number)}\n"
191        out += f"Covalent radius     : {self.get_covalent_radius(atomic_number)}\n"
192        out += f"Van der Waals radius: {self.get_vdw_radius(atomic_number)}\n"
193        return out.rstrip() + "\x1b[0m"
194
195
196
197def append_molecules(molecules):
198    """
199    Append multiple molecules into a single molecule.
200
201    This function takes a list of Molecule objects and returns a new Molecule
202    object that combines all atoms and bonds from the input molecules.
203
204    Arguments:
205        molecules : (list of Molecule) 
206            The molecules to append.
207
208    Returns:
209        Molecule: A new Molecule object containing all atoms and bonds from the input molecules.
210    """
211    if not molecules:
212        raise ValueError("No molecules provided to append.")
213
214    # Create an instance of vtkMoleculeAppend
215    append_filter = vtkMoleculeAppend()
216
217    # Add each molecule's vtkMolecule to the append filter
218    for mol in molecules:
219        append_filter.AddInputData(mol.molecule)
220
221    # Update the filter to generate the combined molecule
222    append_filter.Update()
223
224    # Get the output molecule from the filter
225    combined_vtk_molecule = append_filter.GetOutput()
226
227    # Create a new Molecule object
228    combined_molecule = Molecule()
229
230    # Set the combined vtkMolecule to the new Molecule object
231    combined_molecule.molecule = combined_vtk_molecule
232
233    # Reconfigure the mapper and actor
234    combined_molecule.mapper.SetInputData(combined_vtk_molecule)
235    combined_molecule.actor.SetMapper(combined_molecule.mapper)
236
237    # Optionally, copy rendering settings from the first molecule
238    if molecules:
239        first_mol = molecules[0]
240        if first_mol.mapper.GetRenderAtoms():
241            combined_molecule.use_ball_and_stick()
242        else:
243            combined_molecule.use_space_filling()
244        combined_molecule.set_atom_radius_scale(
245            first_mol.mapper.GetAtomicRadiusScaleFactor()
246        )
247        combined_molecule.set_bond_radius(first_mol.mapper.GetBondRadius())
248
249    return combined_molecule
250
251
252class Atom:
253    """
254    A class representing an atom in a molecule, fully wrapping vtkAtom.
255
256    Provides access to all methods and properties of vtkAtom as documented in:
257    https://vtk.org/doc/nightly/html/classvtkAtom.html
258    """
259
260    def __init__(self, molecule, atom_id):
261        """
262        Initialize the Atom with a reference to the molecule and its ID.
263
264        Arguments:
265            molecule : (vtkMolecule)
266                The molecule containing this atom.
267            atom_id : (int)
268                The ID of the atom in the molecule.
269        """
270        self.molecule = molecule
271        self.atom_id = atom_id
272        self.vtk_atom = self.molecule.GetAtom(self.atom_id)
273
274    def get_atomic_number(self):
275        """Get the atomic number of the atom.
276
277        Returns:
278            The atomic number.
279        """
280        return self.vtk_atom.GetAtomicNumber()
281
282    def set_atomic_number(self, atomic_number):
283        """Set the atomic number of the atom.
284
285        Arguments:
286            atomic_number : (int)
287                The new atomic number.
288        """
289        self.vtk_atom.SetAtomicNumber(atomic_number)
290
291    def get_position(self):
292        """Get the position of the atom as a NumPy array.
293
294        Returns:
295            Array of shape (3,) with [x, y, z] coordinates.
296        """
297        pos = self.vtk_atom.GetPosition()
298        return np.array([pos[0], pos[1], pos[2]])
299
300    def set_position(self, position):
301        """Set the position of the atom.
302
303        Arguments:
304            position : (list or np.ndarray)
305                The new [x, y, z] coordinates.
306        """
307        pos = np.asarray(position, dtype=float)
308        self.vtk_atom.SetPosition(pos[0], pos[1], pos[2])
309
310    def get_atom_id(self):
311        """Get the ID of this atom.
312
313        Returns:
314            The atom's ID within the molecule.
315        """
316        return self.atom_id
317
318    def get_molecule(self):
319        """Get the molecule this atom belongs to.
320
321        Returns:
322            The parent molecule.
323        """
324        return self.molecule
325
326    def __repr__(self):
327        return f"Atom(ID={self.atom_id}, AtomicNumber={self.atomic_number}, Position={self.position})"
328
329
330class Molecule:
331    def __init__(self, pdb_file=None):
332        # Create an empty molecule
333        self.molecule = vtkMolecule()
334
335        # Configure the mapper and actor for rendering
336        self.mapper = vtkMoleculeMapper()
337        self.actor = vtkActor()
338        self.property = self.actor.GetProperty()
339
340        if pdb_file:
341            # Create and configure the PDB reader
342            reader = vtkPDBReader()
343            reader.SetFileName(pdb_file)
344            reader.Update()
345
346            # Get the PDB data
347            pdb_data = reader.GetOutput()
348            # print the point data available
349            # print(pdb_data.GetPointData())
350            # Array 0 name = atom_type
351            # Array 1 name = atom_types
352            # Array 2 name = residue
353            # Array 3 name = chain
354            # Array 4 name = secondary_structures
355            # Array 5 name = secondary_structures_begin
356            # Array 6 name = secondary_structures_end
357            # Array 7 name = ishetatm
358            # Array 8 name = model
359            # Array 9 name = rgb_colors
360            # Array 10 name = radius
361            # for i in range(pdb_data.GetPointData().GetNumberOfArrays()):
362            #     print(pdb_data.GetPointData().GetArray(i))
363            points = pdb_data.GetPoints()
364            point_data = pdb_data.GetPointData()
365
366            # Extract atom information and add to molecule
367            for i in range(points.GetNumberOfPoints()):
368                position = points.GetPoint(i)
369                # Default to Carbon if atomic number is not available
370                atomic_number = 6
371                if point_data.GetScalars("atom_type"):
372                    atomic_number = int(point_data.GetScalars("atom_type").GetValue(i))
373                # Add atom to molecule
374                # if point_data.GetScalars("rgb_colors"):
375                #     color = point_data.GetScalars("rgb_colors").GetTuple3(i)
376                #     self.actor.GetProperty().SetColor(color)
377                self.molecule.AppendAtom(atomic_number, position)
378
379            # Add bonds if available
380            if pdb_data.GetLines():
381                lines = pdb_data.GetLines()
382                lines.InitTraversal()
383                id_list = vtkIdList()
384                while lines.GetNextCell(id_list):
385                    if id_list.GetNumberOfIds() == 2:
386                        self.molecule.AppendBond(
387                            id_list.GetId(0),
388                            id_list.GetId(1),
389                            1,  # Default to single bond
390                        )
391
392        # Set the molecule as input to the mapper
393        self.mapper.SetInputData(self.molecule)
394        self.actor.SetMapper(self.mapper)
395
396        # Apply default rendering style
397        self.use_ball_and_stick()
398
399    def append_atom(self, position=None, atomic_number=6):
400        """Add an atom to the molecule with optional position and atomic number.
401
402        Arguments:
403            position: (list or np.ndarray)
404                [x, y, z] coordinates. Defaults to [0, 0, 0].
405            atomic_number : (int)
406                Atomic number (e.g., 6 for Carbon). Defaults to 6.
407
408        Returns:
409            The added atom object.
410        """
411        if position is None:
412            position = [0, 0, 0]
413        pos = np.asarray(position, dtype=float)
414        vtk_atom = self.molecule.AppendAtom(atomic_number, pos[0], pos[1], pos[2])
415        atom_id = (
416            self.molecule.GetNumberOfAtoms() - 1
417        )  # The ID will be the index of the last added atom
418        return Atom(self.molecule, atom_id)
419
420    def get_atom(self, atom_id):
421        """Retrieve an atom by its ID.
422
423        Arguments:
424            atom_id : (int)
425                The ID of the atom.
426
427        Returns:
428            The atom object.
429        """
430        if atom_id >= self.get_number_of_atoms():
431            raise ValueError(f"Atom ID {atom_id} exceeds number of atoms.")
432        return Atom(self.molecule, atom_id)
433
434    def remove_atom(self, atom_id):
435        """Remove an atom by its ID.
436
437        Arguments:
438            atom_id (int): The ID of the atom to remove.
439        """
440        if atom_id >= self.get_number_of_atoms():
441            raise ValueError(f"Atom ID {atom_id} exceeds number of atoms.")
442        self.molecule.RemoveAtom(atom_id)
443        # Update the mapper to reflect the changes
444        self.mapper.SetInputData(self.molecule)
445        self.mapper.Update()
446        # Update the actor to reflect the changes
447        self.actor.SetMapper(self.mapper)
448        self.actor.Update()
449        # Update the property to reflect the changes
450        self.property.SetColor(1.0, 0.8, 0.6)  # Reset to default color
451        self.property.SetOpacity(1.0)
452
453    def get_array(self, name):
454        """Get a point data array by name.
455
456        The following arrays are available:
457
458        - atom_type: Atomic number of the atom.
459        - atom_types: Atomic type of the atom.
460        - residue: Residue name.
461        - chain: Chain identifier.
462        - secondary_structures: Secondary structure type.
463        - secondary_structures_begin: Start index of the secondary structure.
464        - secondary_structures_end: End index of the secondary structure.
465        - ishetatm: Is the atom a heteroatom?
466        - model: Model number.
467        - rgb_colors: RGB color of the atom.
468        - radius: Radius of the atom.
469
470        Arguments:
471            name : (str)
472                The name of the array.
473
474        Returns:
475            The array data.
476        """
477        if not self.molecule.GetPointData().HasArray(name):
478            raise ValueError(f"Array '{name}' not found in molecule.")
479        return vtk_to_numpy(self.molecule.GetPointData().GetArray(name))
480
481    def append_bond(self, atom1, atom2, order=1):
482        """Add a bond between two atoms.
483
484        Arguments:
485            atom1 : (Atom or int) 
486                The first atom or its ID.
487            atom2 : (Atom or int)
488                The second atom or its ID.
489            order : (int)
490                Bond order (1=single, 2=double, 3=triple).
491        """
492        atom1_id = atom1.atom_id if isinstance(atom1, Atom) else atom1
493        atom2_id = atom2.atom_id if isinstance(atom2, Atom) else atom2
494        self.molecule.AppendBond(atom1_id, atom2_id, order)
495
496    def get_number_of_atoms(self):
497        """Get the number of atoms in the molecule.
498
499        Returns:
500            Number of atoms.
501        """
502        return self.molecule.GetNumberOfAtoms()
503
504    def get_number_of_bonds(self):
505        """Get the number of bonds in the molecule.
506
507        Returns:
508            Number of bonds.
509        """
510        return self.molecule.GetNumberOfBonds()
511
512    def get_bond(self, bond_id):
513        """Get bond information by ID (simplified as VTK bond access is limited).
514
515        Arguments:
516            bond_id : (int)
517                The ID of the bond.
518
519        Returns:
520            (atom1_id, atom2_id, order).
521        """
522        if bond_id >= self.get_number_of_bonds():
523            raise ValueError(f"Bond ID {bond_id} exceeds number of bonds.")
524        bond = self.molecule.GetBond(bond_id)
525        return (bond.GetBeginAtomId(), bond.GetEndAtomId(), bond.GetOrder())
526
527    def get_atom_positions(self):
528        """Get the positions of all atoms.
529
530        Returns:
531            Array of shape (n_atoms, 3) with [x, y, z] coordinates.
532        """
533        n_atoms = self.get_number_of_atoms()
534        positions = np.zeros((n_atoms, 3))
535        for i in range(n_atoms):
536            pos = self.molecule.GetAtom(i).GetPosition()
537            positions[i] = [pos[0], pos[1], pos[2]]
538        return positions
539
540    def set_atom_positions(self, positions):
541        """
542        Set the positions of all atoms.
543
544        Arguments:
545            positions : (np.ndarray)
546                Array of shape (n_atoms, 3) with [x, y, z] coordinates.
547        """
548        n_atoms = self.get_number_of_atoms()
549        positions = np.asarray(positions, dtype=float)
550        if positions.shape != (n_atoms, 3):
551            raise ValueError(
552                f"Expected positions shape ({n_atoms}, 3), got {positions.shape}"
553            )
554        for i in range(n_atoms):
555            self.molecule.GetAtom(i).SetPosition(positions[i])
556
557    def get_atomic_numbers(self):
558        """
559        Get the atomic numbers of all atoms.
560
561        Returns:
562            List of atomic numbers.
563        """
564        return [
565            self.molecule.GetAtom(i).GetAtomicNumber()
566            for i in range(self.get_number_of_atoms())
567        ]
568
569    def set_atomic_numbers(self, atomic_numbers):
570        """
571        Set the atomic numbers of all atoms.
572
573        Arguments:
574            atomic_numbers : (list)
575                List of atomic numbers.
576        """
577        n_atoms = self.get_number_of_atoms()
578        if len(atomic_numbers) != n_atoms:
579            raise ValueError(
580                f"Expected {n_atoms} atomic numbers, got {len(atomic_numbers)}"
581            )
582        for i, num in enumerate(atomic_numbers):
583            self.molecule.GetAtom(i).SetAtomicNumber(num)
584
585    # Rendering customization methods
586    def use_ball_and_stick(self):
587        """Set the molecule to use ball-and-stick representation."""
588        self.mapper.UseBallAndStickSettings()
589        return self
590
591    def use_space_filling(self):
592        """Set the molecule to use space-filling (VDW spheres) representation."""
593        self.mapper.UseVDWSpheresSettings()
594        return self
595
596    def set_atom_radius_scale(self, scale):
597        """Set the scale factor for atom radii.
598
599        Arguments:
600            scale : (float)
601                Scaling factor for atom spheres.
602        """
603        self.mapper.SetAtomicRadiusScaleFactor(scale)
604        return self
605
606    def set_bond_radius(self, radius):
607        """Set the radius of bonds.
608
609        Arguments:
610            radius : (float)
611                Bond radius in world units.
612        """
613        self.mapper.SetBondRadius(radius)
614        return self
615
616
617class Protein:
618    """
619    A Vedo-compatible class for protein ribbon visualization, wrapping vtkProteinRibbonFilter.
620
621    This class generates a ribbon representation of protein structures from PDB files,
622    vtkMolecule objects, or vtkPolyData, and integrates with Vedo's rendering system.
623
624    Attributes:
625        filter (vtkProteinRibbonFilter): The underlying VTK filter for ribbon generation.
626        mapper (vtkPolyDataMapper): Maps the filter's output to renderable data.
627        actor (vtkActor): The VTK actor for rendering the ribbon.
628    """
629
630    def __init__(self, input_data):
631        """
632        Initialize the ProteinRibbon with input data.
633
634        Arguments:
635            input_data : (str or vtkMolecule or vtkPolyData)
636                - Path to a PDB file (str)
637                - A vtkMolecule object
638                - A vtkPolyData object
639
640        Raises:
641            `ValueError` If the input_data type is not supported.
642        """
643
644        # Handle different input types
645        if isinstance(input_data, str):
646            # Read PDB file using vtkPDBReader
647            reader = vtkPDBReader()
648            reader.SetFileName(input_data)
649            reader.Update()
650            self.input_data = reader.GetOutput()
651        elif isinstance(input_data, vtkMolecule):
652            self.input_data = input_data
653        elif isinstance(input_data, vtkPolyData):
654            self.input_data = input_data
655        else:
656            raise ValueError(
657                "Input must be a PDB file path, vtkMolecule, or vtkPolyData."
658            )
659
660        # Create and configure the ribbon filter
661        self.filter = vtkProteinRibbonFilter()
662        self.filter.SetInputData(self.input_data)
663        self.filter.Update()
664
665        # Set up the mapper and actor for rendering
666        self.mapper = vtkPolyDataMapper()
667        self.mapper.SetInputConnection(self.filter.GetOutputPort())
668        self.actor = vtkActor()
669        self.actor.SetMapper(self.mapper)
670        self.property = self.actor.GetProperty()
671
672        # Set default visual properties
673        self.property.SetColor(1.0, 0.8, 0.6)  # Soft peach color
674        self.property.SetOpacity(1.0)
675
676    def set_coil_width(self, width):
677        """
678        Set the width of the coil regions in the ribbon.
679
680        Arguments:
681            width : (float)
682                The width of the coil regions.
683        """
684        self.filter.SetCoilWidth(width)
685        self.filter.Update()
686        return self
687
688    def set_helix_width(self, width):
689        """
690        Set the width of the helix regions in the ribbon.
691
692        Arguments:
693            width : (float)
694                The width of the helix regions.
695        """
696        self.filter.SetHelixWidth(width)
697        self.filter.Update()
698        return self
699
700    def set_sphere_resolution(self, resolution):
701        """
702        Set the resolution of spheres used in the ribbon representation.
703
704        Arguments:
705            resolution : (int)
706                The resolution of the spheres.
707        """
708        self.filter.SetSphereResolution(resolution)
709        self.filter.Update()
710        return self