fennol.models.physics.polarisation

  1#!/usr/bin/env python3
  2"""Polarization model for FENNOL.
  3
  4Created by C. Cattin 2024
  5"""
  6
  7import flax.linen as nn
  8import jax
  9import jax.numpy as jnp
 10from typing import Optional, ClassVar
 11
 12from fennol.utils import AtomicUnits as Au
 13
 14
 15class Polarization(nn.Module):
 16    """Polarization model with Thole damping scheme.
 17    
 18    FID: POLARIZATION
 19    """
 20
 21    energy_key: Optional[str] = None
 22    """Key of the energy in the outputs."""
 23    graph_key: str = 'graph'
 24    """Key of the graph in the inputs."""
 25    polarizability_key: str = 'polarizability'
 26    """Key of the polarizability in the inputs."""
 27    electric_field_key: str = 'electric_field'
 28    """Key of the electric field in the inputs."""
 29    induced_dipoles_key: str = 'induced_dipoles'
 30    """Key of the induced dipole in the outputs."""
 31    damping_param_mutual: float = 0.39
 32    """Damping parameter for mutual polarization."""
 33    neglect_mutual: bool = False
 34    """Neglect the mutual polarization term (like in iAMOEBA)."""
 35    _energy_unit: str = 'Ha'
 36    """The energy unit of the model. **Automatically set by FENNIX**"""
 37
 38    FID: ClassVar[str] = 'POLARIZATION'
 39
 40    @nn.compact
 41    def __call__(self, inputs):
 42        """Forward pass of the polarization model.
 43
 44        Parameters
 45        ----------
 46        inputs : dict
 47            Input dictionary containing all the info about the system.
 48            This dictionary is given from the FENNIX class.
 49        """
 50        # Species
 51        species = inputs['species']
 52        graph = inputs[self.graph_key]
 53        # Graph
 54        edge_src, edge_dst = graph['edge_src'], graph['edge_dst']
 55        # Distances and vector between each pair of atoms in atomic units
 56        distances = graph['distances']
 57        rij = distances / Au.BOHR
 58        vec_ij = graph['vec'] / Au.BOHR
 59        # Polarizability
 60        polarizability = (
 61            inputs[self.polarizability_key] / Au.BOHR**3
 62        )
 63
 64        pol_src = polarizability[edge_src]
 65        pol_dst = polarizability[edge_dst]
 66        alpha_ij = pol_dst * pol_src
 67
 68        # The output is a dictionary with the polarization energy
 69        output = {}
 70
 71        ######################
 72        # Interaction matrix #
 73        ######################
 74        # Diagonal terms
 75        tii = 1 / polarizability[:, None]
 76        
 77
 78        ##################
 79        # Electric field #
 80        ##################
 81        electric_field = inputs[self.electric_field_key]
 82
 83        ###############################
 84        # Electric point dipole moment#
 85        ###############################
 86        if self.neglect_mutual:
 87            electric_field = electric_field.reshape(-1, 3)
 88            mu = polarizability[:,None]*electric_field
 89            mu_ = jax.lax.stop_gradient(mu)
 90            tmu = tii*mu_
 91            
 92        else:
 93            # Effective distance
 94            uij = rij / alpha_ij ** (1 / 6)
 95            # Damping terms
 96            exp = jnp.exp(-self.damping_param_mutual * uij**3)
 97            lambda_3 = 1 - exp
 98            lambda_5 = 1 - (1 + self.damping_param_mutual * uij**3) * exp
 99            # Off-diagonal terms
100            tij = (
101                3 * lambda_5[:, None, None]
102                    * vec_ij[:, :, None] * vec_ij[:, None, :]
103                    / rij[:, None, None]**5
104                - jnp.eye(3)[None, :, :] * lambda_3[:, None, None]
105                    / rij[:, None, None]**3
106            )
107
108            def matvec(mui):
109                """Compute the matrix vector product of T and mu."""
110                mui = mui.reshape(-1, 3)
111                tmu_self = tii * mui
112                tmu_pair = jnp.einsum("jab,jb->ja", tij, mui[edge_dst])
113                tmu = (
114                    jax.ops.segment_sum(tmu_pair, edge_src, species.shape[0])
115                    + tmu_self
116                )
117                return tmu.flatten()
118            mu = jax.scipy.sparse.linalg.cg(matvec, electric_field)[0]
119            mu_ = jax.lax.stop_gradient(mu)
120
121            # Matrix vector product
122            tmu = matvec(mu_)
123
124        # Polarization energy
125        pol_energy = (
126            (0.5 * tmu - electric_field) * mu_
127        ).reshape(-1, 3).sum(axis=1)
128
129        # Output
130        output[self.electric_field_key] = electric_field.reshape(-1, 3)
131        output[self.induced_dipoles_key] = mu.reshape(-1, 3) * Au.BOHR
132        energy_key = (
133            self.energy_key if self.energy_key is not None else self.name
134        )
135        energy_unit = Au.get_multiplier(self._energy_unit)
136        output[energy_key] = pol_energy*energy_unit
137        output['tmu'] = tmu.reshape(-1, 3)
138
139        return {**inputs, **output}
140
141
142if __name__ == "__main__":
143    pass