Esempio n. 1
0
def test_shader_callback():

    cone = ConeSource()
    coneMapper = PolyDataMapper()
    coneMapper.SetInputConnection(cone.GetOutputPort())
    actor = Actor()
    actor.SetMapper(coneMapper)

    test_values = []

    def callbackLow(_caller, _event, calldata=None):
        program = calldata
        if program is not None:
            test_values.append(0)

    id_observer = fs.add_shader_callback(actor, callbackLow, 0)

    with pytest.raises(Exception):
        fs.add_shader_callback(actor, callbackLow, priority='str')

    mapper = actor.GetMapper()
    mapper.RemoveObserver(id_observer)

    scene = window.Scene()
    scene.add(actor)

    arr1 = window.snapshot(scene, size=(200, 200))
    assert len(test_values) == 0

    test_values = []

    def callbackHigh(_caller, _event, calldata=None):
        program = calldata
        if program is not None:
            test_values.append(999)

    def callbackMean(_caller, _event, calldata=None):
        program = calldata
        if program is not None:
            test_values.append(500)

    fs.add_shader_callback(actor, callbackHigh, 999)
    fs.add_shader_callback(actor, callbackLow, 0)

    id_mean = fs.add_shader_callback(actor, callbackMean, 500)

    # check the priority of each call
    arr2 = window.snapshot(scene, size=(200, 200))
    assert np.abs(
        [test_values[0] - 999, test_values[1] - 500,
         test_values[2] - 0]).sum() == 0

    # check if the correct observer was removed
    mapper.RemoveObserver(id_mean)
    test_values = []

    arr3 = window.snapshot(scene, size=(200, 200))
    assert np.abs([test_values[0] - 999, test_values[1] - 0]).sum() == 0
Esempio n. 2
0
def ribbon(molecule):
    """Create an actor for ribbon molecular representation.

    Parameters
    ----------
    molecule : Molecule
        The molecule to be rendered.

    Returns
    -------
    molecule_actor : vtkActor
        Actor created to render the rubbon representation of the molecule to be
        visualized.

    References
    ----------
    Richardson, J.S. The anatomy and taxonomy of protein structure
    `Advances in Protein Chemistry, 1981, 34, 167-339.
    <https://doi.org/10.1016/S0065-3233(08)60520-3>`_
    """
    coords = get_all_atomic_positions(molecule)
    all_atomic_numbers = get_all_atomic_numbers(molecule)
    num_total_atoms = molecule.total_num_atoms
    secondary_structures = np.ones(num_total_atoms)
    for i in range(num_total_atoms):
        secondary_structures[i] = ord('c')
        resi = molecule.residue_seq[i]
        for j, _ in enumerate(molecule.sheet):
            sheet = molecule.sheet[j]
            if molecule.chain[i] != sheet[0] or resi < sheet[1] or \
               resi > sheet[3]:
                continue
            secondary_structures[i] = ord('s')

        for j, _ in enumerate(molecule.helix):
            helix = molecule.helix[j]
            if molecule.chain[i] != helix[0] or resi < helix[1] or \
               resi > helix[3]:
                continue
            secondary_structures[i] = ord('h')

    output = PolyData()

    # for atomic numbers
    atomic_num_arr = nps.numpy_to_vtk(num_array=all_atomic_numbers,
                                      deep=True,
                                      array_type=VTK_ID_TYPE)

    # setting the array name to atom_type as vtkProteinRibbonFilter requires
    # the array to be named atom_type
    atomic_num_arr.SetName("atom_type")

    output.GetPointData().AddArray(atomic_num_arr)

    # for atom names
    atom_names = StringArray()

    # setting the array name to atom_types as vtkProteinRibbonFilter requires
    # the array to be named atom_types
    atom_names.SetName("atom_types")
    atom_names.SetNumberOfTuples(num_total_atoms)
    for i in range(num_total_atoms):
        atom_names.SetValue(i, molecule.atom_names[i])

    output.GetPointData().AddArray(atom_names)

    # for residue sequences
    residue_seq = nps.numpy_to_vtk(num_array=molecule.residue_seq,
                                   deep=True,
                                   array_type=VTK_ID_TYPE)
    residue_seq.SetName("residue")
    output.GetPointData().AddArray(residue_seq)

    # for chain
    chain = nps.numpy_to_vtk(num_array=molecule.chain,
                             deep=True,
                             array_type=VTK_UNSIGNED_CHAR)
    chain.SetName("chain")
    output.GetPointData().AddArray(chain)

    # for secondary structures
    s_s = nps.numpy_to_vtk(num_array=secondary_structures,
                           deep=True,
                           array_type=VTK_UNSIGNED_CHAR)
    s_s.SetName("secondary_structures")
    output.GetPointData().AddArray(s_s)

    # for secondary structures begin
    newarr = np.ones(num_total_atoms)
    s_sb = nps.numpy_to_vtk(num_array=newarr,
                            deep=True,
                            array_type=VTK_UNSIGNED_CHAR)
    s_sb.SetName("secondary_structures_begin")
    output.GetPointData().AddArray(s_sb)

    # for secondary structures end
    newarr = np.ones(num_total_atoms)
    s_se = nps.numpy_to_vtk(num_array=newarr,
                            deep=True,
                            array_type=VTK_UNSIGNED_CHAR)
    s_se.SetName("secondary_structures_end")
    output.GetPointData().AddArray(s_se)

    # for is_hetatm
    is_hetatm = nps.numpy_to_vtk(num_array=molecule.is_hetatm,
                                 deep=True,
                                 array_type=VTK_UNSIGNED_CHAR)
    is_hetatm.SetName("ishetatm")
    output.GetPointData().AddArray(is_hetatm)

    # for model
    model = nps.numpy_to_vtk(num_array=molecule.model,
                             deep=True,
                             array_type=VTK_UNSIGNED_INT)
    model.SetName("model")
    output.GetPointData().AddArray(model)

    table = PTable()

    # for colors and radii of hetero-atoms
    radii = np.ones((num_total_atoms, 3))
    rgb = np.ones((num_total_atoms, 3))

    for i in range(num_total_atoms):
        radii[i] = np.repeat(table.atomic_radius(all_atomic_numbers[i], 'VDW'),
                             3)
        rgb[i] = table.atom_color(all_atomic_numbers[i])

    Rgb = nps.numpy_to_vtk(num_array=rgb,
                           deep=True,
                           array_type=VTK_UNSIGNED_CHAR)
    Rgb.SetName("rgb_colors")
    output.GetPointData().SetScalars(Rgb)

    Radii = nps.numpy_to_vtk(num_array=radii, deep=True, array_type=VTK_FLOAT)
    Radii.SetName("radius")
    output.GetPointData().SetVectors(Radii)

    # setting the coordinates
    points = numpy_to_vtk_points(coords)
    output.SetPoints(points)

    ribbonFilter = ProteinRibbonFilter()
    ribbonFilter.SetInputData(output)
    ribbonFilter.SetCoilWidth(0.2)
    ribbonFilter.SetDrawSmallMoleculesAsSpheres(0)
    mapper = PolyDataMapper()
    mapper.SetInputConnection(ribbonFilter.GetOutputPort())
    molecule_actor = Actor()
    molecule_actor.SetMapper(mapper)
    return molecule_actor