Constraints

Constraints in tmol allow you to define custom score functions in Python, alongside sets of atoms that the this score function will be evaluated on.

Usage

Using constraints requires:

1. A constraint function. There are several common functions built-in that you may use: harmonic, bounded, and circularharmonic.
2. A integer Tensor with the indices of each atom in the constraint. The shape of this will be [n_constraints, n_atoms, 3], where the 3 values in the last dimension represent the [pose_index, residue_index, atom_index] for the atom in question.
3. A float Tensor containing the parameters for the function for each constraint. The shape of this will be [n_constraints, n_parameters_per_constraint].

    # get the constraint set for the post (this will create one if it does not exist)
    constraints = pose_stack.get_constraint_set()

    # create tensors to hold the atom indices of each constraint
    cnstr_atoms = torch.full((1, 2, 3), 0, dtype=torch.int32, device=torch_device)
    # ... and the parameters for each constraint (in this case, just 1 distance value per constraint)
    cnstr_params = torch.full((1, 1), 0, dtype=torch.float32, device=torch_device)

    # now we fill those tensors for each constraint:

    # grab the block type for each residue that we want the atoms from
    res1_type = pose_stack.block_type(0, 3)
    res2_type = pose_stack.block_type(0, 4)
    # get the [pose, residue, atom] indices for each atom in our constraint
    cnstr_atoms[0, 0] = torch.tensor([0, 3, res1_type.atom_to_idx["C"]])
    cnstr_atoms[0, 1] = torch.tensor([0, 4, res2_type.atom_to_idx["N"]])

    # now fill the distance parameter for this constraint
    cnstr_params[0, 0] = 1.47

    # add the constraints to the constraint set, using the built-in harmonic function
    constraints.add_constraints(
        ConstraintEnergyTerm.harmonic, cnstr_atoms, cnstr_params
    )

Constraint Functions

In addition to the built-in functions, you may define your own. These functions must match the following rules:

1. The function must take 2 arguments.
2. The first argument must be a Tensor of atom coordinates for each set of atoms. The shape of this tensor must be [n_constraints x n_atoms x dimension]. For example, if your function is being applied to 5 different sets of atoms, with 4 atoms per set, this would be a float tensor of 5x4x3.
3. The second argument is for parameters to your function. This argument is optional, but must be included in your function signature even if it is unused. If you use this, the parameters must be a Tensor of float values, with a shape of [n_constraints x n_parameters_per_constraint]. An example use of this would be to individually describe distances in an atom pair constraint. If you had 10 atom pair constraints and a single distance parameter for each of them, this would be a 10x1 tensor.
4. The return value must be a float Tensor of shape [n_constraints], with your evaluated score for each constraint at each index.
5. Your function must take at least 1 and at most 4 atoms.
6. During score attribution, the constraint score term will attribute half the score to the residue of the first atom in the atom set, and half to the last. Both atoms may be from the same residue, which will attribute that residue the full score.

As an example, here is the code for an atom-pair distance constraint:

    def harmonic(atoms, params):
        atoms1 = atoms[:, 0]
        atoms2 = atoms[:, 1]
        dist = torch.linalg.norm(atoms1 - atoms2, dim=-1)
        return (dist - params[:, 0]) ** 2