Source code for mythos.energy.dna2.stacking

"""Stacking energy function for DNA1 model."""

import chex
import jax.numpy as jnp
from jax import vmap
from typing_extensions import override

import mythos.energy.dna1 as jd_energy1
import mythos.energy.dna2.nucleotide as dna2_nucleotide
import mythos.utils.types as typ




@chex.dataclass(frozen=True)
class Stacking(jd_energy1.Stacking):
    """Stacking energy function for DNA2 model."""

    params: jd_energy1.StackingConfiguration



    def pairwise_energies(
        self,
        body: dna2_nucleotide.Nucleotide,
        seq: typ.Discrete_Sequence,
        bonded_neighbors: typ.Arr_Bonded_Neighbors_2,
    ) -> typ.Arr_Bonded_Neighbors:
        """Computes the stacking energy for each bonded pair."""
        # Compute sequence-independent energy for each bonded pair
        v_stack = self.compute_v_stack(
            body.stack_sites, body.back_sites_dna1, body.base_normals, body.cross_prods, bonded_neighbors
        )

        # Compute sequence-dependent weight for each bonded pair
        nn_i = bonded_neighbors[:, 0]
        nn_j = bonded_neighbors[:, 1]

        if self.params.pseq:
            stack_weights = vmap(self.pseq_weights, (0, 0, None))(nn_i, nn_j, self.params.pseq)
        else:
            stack_weights = self.params.eps_stack[seq[nn_i], seq[nn_j]]
        return jnp.multiply(stack_weights, v_stack)




    @override
    def compute_energy(self, nucleotide: dna2_nucleotide.Nucleotide) -> typ.Scalar:
        dgs = self.pairwise_energies(nucleotide, self.seq, self.bonded_neighbors)
        return dgs.sum()