fennol.models.embeddings.charge_embeddings

 1import jax
 2import jax.numpy as jnp
 3import flax.linen as nn
 4from typing import Sequence, Dict, Union, ClassVar, Optional
 5import numpy as np
 6from ...utils.periodic_table import (
 7    VALENCE_ELECTRONS,
 8)
 9
10class ChargeHypothesis(nn.Module):
11    """Embedding with total charge constraint vi Multiple Neural Charge Equilibration.
12
13    FID: CHARGE_HYPOTHESIS
14
15    """
16
17    embedding_key: str
18    """Key for the embedding in the inputs that is used to predict an 'electron affinity' weight"""
19    output_key: Union[str, None] = None
20    """Key for the charges in the outputs"""
21    total_charge_key: str = "total_charge"
22    """Key for the total charge in the inputs"""
23    ncharges: int = 10
24    """The number of charge hypothesis"""
25    mode: str = "qeq"
26    """Charge distribution mode. Only 'qeq' available for now."""
27    squeeze: bool = True
28
29    FID: ClassVar[str] = "CHARGE_HYPOTHESIS"
30
31    @nn.compact
32    def __call__(self, inputs):
33        embedding = inputs[self.embedding_key].astype(inputs["coordinates"].dtype)
34
35        wi = jax.nn.softplus(
36            nn.Dense(self.ncharges, use_bias=True, name="wi")(embedding)
37        )
38
39        batch_index = inputs["batch_index"]
40        nsys = inputs["natoms"].shape[0]
41        wtot = jax.ops.segment_sum(wi, batch_index, nsys)
42
43        Qtot = (
44            inputs[self.total_charge_key].astype(wi.dtype)
45            if self.total_charge_key in inputs
46            else jnp.zeros(nsys, dtype=wi.dtype)
47        )
48        if Qtot.ndim == 0:
49            Qtot = Qtot * jnp.ones(nsys, dtype=wi.dtype)
50
51        
52        qtilde = nn.Dense(self.ncharges, use_bias=True, name="qi")(embedding)
53        qtot = jax.ops.segment_sum(qtilde, batch_index, nsys)
54        dq = Qtot[:, None] - qtot
55        f = (dq / wtot)[batch_index]
56
57        if "alch_group" in inputs:
58            alch_group = inputs["alch_group"]
59            lambda_e = inputs["alch_elambda"]
60            Qligand = inputs.get("alch_ligand_charge", jnp.zeros(nsys))
61            if Qligand.ndim == 0:
62                Qligand = Qligand * jnp.ones(nsys, dtype=wi.dtype)
63
64            alch_index = alch_group + 2 * batch_index
65            qtot = jax.ops.segment_sum(qtilde, alch_index, 2 * nsys).reshape(
66                nsys, 2, self.ncharges
67            )
68            wtot = jax.ops.segment_sum(wi, alch_index, 2 * nsys).reshape(
69                nsys, 2, self.ncharges
70            )
71            # compute ligand alone
72            dq = Qligand[:, None] - qtot[:, 1, :]
73            f1 = (dq / wtot[:, 1, :])[batch_index]
74            # q1 = alch_group*(qtilde + wi * f1[batch_index])
75
76            # qtot = jax.ops.segment_sum(qtilde*(1-alch_group), batch_index, nsys)
77            # wtot = jax.ops.segment_sum(wi*(1-alch_group), batch_index, nsys)
78            # compute solvent alone
79            dq = (Qtot - Qligand)[:, None] - qtot[:, 0, :]
80            f0 = (dq / wtot[:, 0, :])[batch_index]
81            # q0 = (1-alch_group)*(qtilde + wi * f0[batch_index])
82
83            # alch_q = alch_group*q1 + (1-alch_group)*q0
84            # q = lambda_e*q + (1-lambda_e)*alch_q
85
86            alch_f = (1 - alch_group[:, None]) * f0 + alch_group[:, None] * f1
87            f = lambda_e * f + (1 - lambda_e) * alch_f
88
89        q = qtilde + wi * f
90
91        if self.squeeze and self.ncharges == 1:
92            q = jnp.squeeze(q, axis=-1)
93
94        output_key = self.output_key if self.output_key is not None else self.name
95        return {
96            **inputs,
97            output_key: q,
98        }