Class TorsionalDiffusion

diffpack/task.py:19–220

NewTorsionalDiffusion is a class for simulating the torsional diffusion of a protein model. It inherits from the tasks.Task and core.Configurable classes and uses these to setup and control the diffusion simulation. Attributes: eps (float): A small number to avoid division

@R.register("tasks.TorsionalDiffusion")
class TorsionalDiffusion(tasks.Task, core.Configurable):
    """
    NewTorsionalDiffusion is a class for simulating the torsional diffusion of a protein model.

    It inherits from the tasks.Task and core.Configurable classes and uses these to setup and control the diffusion
    simulation.

    Attributes:
        eps (float): A small number to avoid division by zero errors.
        _option_members (set): A set containing the names of class attributes.
        model (nn.Module): The neural network model to be used.
        schedule_1pi_periodic (SO2VESchedule): The schedule for 1pi periodic tasks.
        schedule_2pi_periodic (SO2VESchedule): The schedule for 2pi periodic tasks.
        num_mlp_layer (int): The number of layers in the model.
        graph_construction_model (Optional[Any]): The model used for graph construction.
        verbose (int): Verbosity level.
        train_chi_id (Optional[Any]): Chi angle for training ranging from 0 to 3. If not specified, random chi angles are trained.
    """
    NUM_CHI_ANGLES = 4
    eps = 1e-10
    _option_members = {"task", "criterion", "metric"}

    def __init__(self, sigma_embedding: nn.Module,
                 model: nn.Module,
                 torsion_mlp_hidden_dims: list,
                 schedule_1pi_periodic: SO2VESchedule,
                 schedule_2pi_periodic: SO2VESchedule,
                 graph_construction_model: Optional[Any] = None,
                 verbose: int = 0,
                 train_chi_id: Optional[Any] = None, ):
        super(TorsionalDiffusion, self).__init__()
        self.torsion_mlp_hidden_dims = torsion_mlp_hidden_dims
        self.model_list = nn.ModuleList([deepcopy(model) for _ in range(self.NUM_CHI_ANGLES)])
        self.sigma_embedding_list = nn.ModuleList([deepcopy(sigma_embedding) for _ in range(self.NUM_CHI_ANGLES)])
        self.torsion_mlp_list = nn.ModuleList([layers.MLP(self.model_list[i].output_dim, torsion_mlp_hidden_dims
                                                          + [4,]) for i in range(self.NUM_CHI_ANGLES)])
        self.schedule_2pi_periodic = schedule_2pi_periodic
        self.schedule_1pi_periodic = schedule_1pi_periodic
        self.graph_construction_model = graph_construction_model
        self.verbose = verbose
        self.train_chi_id = train_chi_id

    def forward(self, batch):
        all_loss = torch.tensor(0, dtype=torch.float32, device=self.device)
        metric = {}

        # Sample from the schedule
        protein = batch['graph']
        t = self.schedule_1pi_periodic.sample_train_t(shape=(protein.batch_size,)).to(self.device)

        # Add noise to protein
        train_chi_id = np.random.randint(self.NUM_CHI_ANGLES) if self.train_chi_id is None else self.train_chi_id
        batch = self.add_noise(batch, t, chi_id=train_chi_id)

        # Predict and take loss
        pred = self.predict(batch, all_loss, metric)
        target = self.target(batch)