Dipole Model

The dipole model1 requires the same dictionary as input as the potential model. The only difference is the model_params that can be set. They are listed below along with their default values.

Parameter Default Description
max_dipole False When set to a number, exclude dipoles above this value in loss function
d_scale 1 Dipole unit during training
use_d_per_atom False Use the per-atom dipole in place of total dipole in loss function
log_d_per_atom True Log the per-atom dipole error, automatically enabled with use_d_per_atoms=True
use_d_weight False Scale the energy loss according to the 'd_weight' Tensor in the dataset
use_l2 False Include L2 regularization in loss
d_loss_multiplier 1 Weight of dipole loss
l2_loss_multiplier 1 Weight of l2

Variants

Since the PiNet2 update several new variants of the dipole model have become available.

Usage

Apart from the models containing the atomic dipole (AD), all models are compatible with both PiNet and PiNet2.

These models can be selected by changing name of the model in model_params (See overview for further instructions.).

Parameters

Parameter Default Description
max_dipole False When set to a number, exclude dipoles above this value in loss function
d_scale 1 Dipole unit during training
d_unit 1 Output unit of dipole during prediction
vector_dipole False Toggle whether to use scalar or vector dipole predictions, should be the same as dipole moment of the dataset
charge_neutrality True Enable charge neutrality
neutral_unit 'system' If charge neutrality is applied, set the charge neutral unit. Choose 'system' for system neutrality, or 'water_molecule' for neutrality per water molecule
regularization True Enable regularization of the interaction prediction, only available for models containing the 'R' term
use_d_per_atom False Use the per-atom dipole in place of total dipole in loss function
log_d_per_atom True Log the per-atom dipole error, automatically enabled with use_d_per_atoms=True
use_d_weight False Scale the energy loss according to the 'd_weight' Tensor in the dataset
use_l2 False Include L2 regularization in loss
d_loss_multiplier 1 Weight of dipole loss
l2_loss_multiplier 1 Weight of l2

Model specifications

pinn.models.AC.AC_dipole_model

The atomic charge (AC) dipole model constructs the dipole moment from atomic charge predictions:

\[ \begin{aligned} \mu = \sum_{i}{}^{1}\mathbb{P}_{i} \cdot \mathbf{r}_{i} \end{aligned} \]
Source code in pinn/models/AC.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
@export_model
def AC_dipole_model(features, labels, mode, params):
    r"""The atomic charge (AC) dipole model constructs the dipole moment from 
    atomic charge predictions:

    $$
    \begin{aligned}
    \mu = \sum_{i}{}^{1}\mathbb{P}_{i} \cdot \mathbf{r}_{i}
    \end{aligned}
    $$
    """

    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)
    p1 = network(features)

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    if model_params['charge_neutrality'] == True:
      if model_params['neutral_unit'] == 'system':
          q_molecule = tf.math.unsorted_segment_sum(p1, ind1[:, 0], nbatch)
          N = tf.math.unsorted_segment_sum(tf.ones_like(ind1, tf.float32), ind1, tf.reduce_max(ind1)+1)

          p_charge = q_molecule/N
          charge_corr = tf.gather(p_charge, ind1)[:,0]
          p1 =  p1 - charge_corr

      if model_params['neutral_unit'] == 'water_molecule':
          q_molecule = tf.math.reduce_sum(tf.reshape(p1,[-1,3]),axis=1)

          p_charge = q_molecule/3
          charge_corr = tf.reshape(tf.stack([p_charge, p_charge, p_charge], axis=1), [1,-1])[0,:]
          p1 =  p1 - charge_corr

    q_tot = tf.math.unsorted_segment_sum(p1, ind1[:, 0], nbatch)
    p1 = tf.expand_dims(p1, axis=1)

    q_d = p1 * features['coord']
    dipole = tf.math.unsorted_segment_sum(q_d, ind1[:, 0], nbatch)

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, q_tot, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, q_tot, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole,
            'charge': q_tot,
            'charges': tf.expand_dims(p1, 0)
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

pinn.models.AD.AD_dipole_model

The atomic dipole (AD) dipole model constructs the dipole moment from atomic dipole predictions:

\[ \begin{aligned} \mu = \sum_{i}{}^{3}\mathbb{P}_{i} \end{aligned} \]
Source code in pinn/models/AD.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@export_model
def AD_dipole_model(features, labels, mode, params):
    r"""The atomic dipole (AD) dipole model constructs the dipole moment from 
    atomic dipole predictions:

    $$
    \begin{aligned}
    \mu = \sum_{i}{}^{3}\mathbb{P}_{i}
    \end{aligned}
    $$
    """
    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)
    p1, output_p3 = network(features)
    p3 = output_p3['p3']
    p3 = tf.squeeze(p3, axis=-1)

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    atomic_d = tf.math.unsorted_segment_sum(p3, ind1[:, 0], nbatch)

    dipole = atomic_d

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole
            #'atomic_d': tf.expand_dims(p3, 0)
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

pinn.models.BC_R.BC_R_dipole_model

The bond charge dipole model with regularization (BC(R)) constructs the dipole moment from atomic pairwise interactions:

\[ \begin{aligned} \mu = \sum_{ij}{}^{1}\mathbb{I}_{ij} \cdot \mathbf{r}_{ij} \end{aligned} \]

L2-regularization is applied to the atomic pairwise interactions.

Source code in pinn/models/BC_R.py 
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@export_model
def BC_R_dipole_model(features, labels, mode, params):
    r"""The bond charge dipole model with regularization (BC(R))
    constructs the dipole moment from atomic pairwise interactions:

    $$
    \begin{aligned}
    \mu = \sum_{ij}{}^{1}\mathbb{I}_{ij} \cdot \mathbf{r}_{ij}
    \end{aligned}
    $$

    L2-regularization is applied to the atomic pairwise interactions.
    """
    if params['network']['name'] == "PiNet":
        params['network']['params'].update({'out_prop':0, 'out_inter':1})

    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)

    if params['network']['name'] == "PiNet2":
        ppred, output_dict = network(features)
        ppred = tf.expand_dims(ppred, axis=1)

        i1 = output_dict['i1']
        i3 = output_dict['i3']

        ipred =  tf.einsum("ixr,ixr->ir", i3, i3) + i1

    else:
        ppred, ipred = network(features)

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    # Compute bond vector
    disp_r = features['diff']

    # Compute atomic dipole
    atomic_d_pairwise = ipred * disp_r
    atomic_d = tf.math.unsorted_segment_sum(atomic_d_pairwise, ind2[:, 0], natoms) 
    dipole = tf.math.unsorted_segment_sum(atomic_d, ind1[:, 0], nbatch)

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, ipred, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, ipred, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole,
            #'atomic_d': tf.expand_dims(atomic_d, 0)
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

pinn.models.AD_OS.AD_OS_dipole_model

The atomic dipole with oxidation state term (AD(OS)) dipole model constructs the dipole moment from atomic dipole predictions and oxidation state charges:

\[ \begin{aligned} \mu = \sum_{i}{}^{3}\mathbb{P}_{i} + q^{ox}_{i} \cdot \mathbf{r}_{i} \end{aligned} \]
Source code in pinn/models/AD_OS.py 
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@export_model
def AD_OS_dipole_model(features, labels, mode, params):
    r"""The atomic dipole with oxidation state term (AD(OS)) dipole model 
    constructs the dipole moment from atomic dipole predictions and oxidation
    state charges:

    $$
    \begin{aligned}
    \mu = \sum_{i}{}^{3}\mathbb{P}_{i} + q^{ox}_{i} \cdot \mathbf{r}_{i}
    \end{aligned}
    $$
    """
    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)
    p1, output_p3 = network(features)
    p3 = output_p3['p3']
    p3 = tf.squeeze(p3, axis=-1)

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    ox = features['oxidation']
    ox = tf.expand_dims(ox, axis=1)

    q_tot = tf.math.unsorted_segment_sum(ox, ind1[:, 0], nbatch)

    q_d = ox * features['coord']
    mol_q_d = tf.math.unsorted_segment_sum(q_d, ind1[:, 0], nbatch)

    atomic_d = tf.math.unsorted_segment_sum(p3, ind1[:, 0], nbatch)

    a_dipole = q_d + p3
    dipole = mol_q_d + atomic_d

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, q_tot, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, q_tot, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole,
            #'charge': q_tot
            #'atomic_d': tf.expand_dims(a_dipole, 0)
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

pinn.models.AC_AD.AC_AD_dipole_model

The AC+AD dipole model constructs the dipole moment from atomic dipole predictions and atomic charge predictions:

\[ \begin{aligned} \mu = \sum_{i}{}^{3}\mathbb{P}_{i} + {}^{1}\mathbb{P}_{i} \cdot \mathbf{r}_{i} \end{aligned} \]
Source code in pinn/models/AC_AD.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@export_model
def AC_AD_dipole_model(features, labels, mode, params):
    r"""The AC+AD  dipole model constructs the dipole moment from 
    atomic dipole predictions and atomic charge predictions:

    $$
    \begin{aligned}
    \mu = \sum_{i}{}^{3}\mathbb{P}_{i} + {}^{1}\mathbb{P}_{i} \cdot \mathbf{r}_{i}    
    \end{aligned}
    $$
    """
    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)
    p1, output_p3 = network(features)
    p3 = output_p3['p3']
    p3 = tf.squeeze(p3, axis=-1)

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    if model_params['charge_neutrality'] == True:
      if model_params['neutral_unit'] == 'system':
          q_molecule = tf.math.unsorted_segment_sum(p1, ind1[:, 0], nbatch)
          N = tf.math.unsorted_segment_sum(tf.ones_like(ind1, tf.float32), ind1, tf.reduce_max(ind1)+1)

          p_charge = q_molecule/N
          charge_corr = tf.gather(p_charge, ind1)[:,0]
          p1 =  p1 - charge_corr

      if model_params['neutral_unit'] == 'water_molecule':
          q_molecule = tf.math.reduce_sum(tf.reshape(p1,[-1,3]),axis=1)

          p_charge = q_molecule/3
          charge_corr = tf.reshape(tf.stack([p_charge, p_charge, p_charge], axis=1), [1,-1])[0,:]
          p1 =  p1 - charge_corr

    q_tot = tf.math.unsorted_segment_sum(p1, ind1[:, 0], nbatch)
    p1 = tf.expand_dims(p1, axis=1)

    q_d = p1 * features['coord']
    mol_q_d = tf.math.unsorted_segment_sum(q_d, ind1[:, 0], nbatch)

    atomic_d = tf.math.unsorted_segment_sum(p3, ind1[:, 0], nbatch)

    a_dipole = q_d + p3
    dipole = mol_q_d + atomic_d

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, q_tot, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, q_tot, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole,
            'charge': q_tot
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

pinn.models.AC_BC_R.AC_BC_R_dipole_model

The AC+BC(R) constructs the dipole moment from atomic charge predictions and atomic pairwise interactions:

\[ \begin{aligned} \mu = \sum_{i}{}^{1}\mathbb{P}_{i} \cdot \mathbf{r}_{i} + \sum_{ij}{}^{1}\mathbb{I}_{ij} \cdot \mathbf{r}_{ij} \end{aligned} \]

L2-regularization is applied to the atomic pairwise interactions.

Source code in pinn/models/AC_BC_R.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@export_model
def AC_BC_R_dipole_model(features, labels, mode, params):
    r"""The AC+BC(R) constructs the dipole moment from 
    atomic charge predictions and atomic pairwise interactions:

    $$
    \begin{aligned}
    \mu = \sum_{i}{}^{1}\mathbb{P}_{i} \cdot \mathbf{r}_{i} + \sum_{ij}{}^{1}\mathbb{I}_{ij} \cdot \mathbf{r}_{ij}
    \end{aligned}
    $$

    L2-regularization is applied to the atomic pairwise interactions.
    """
    if params['network']['name'] == "PiNet":
        params['network']['params'].update({'out_prop':1, 'out_inter':1})

    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)

    if params['network']['name'] == "PiNet2":
        p1, output_dict = network(features)

        i1 = output_dict['i1']
        i3 = output_dict['i3']

        ipred =  tf.einsum("ixr,ixr->ir", i3, i3) + i1

    else:
        p1, ipred = network(features)

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    if model_params['charge_neutrality'] == True:
        if model_params['neutral_unit'] == 'system':
            q_molecule = tf.math.unsorted_segment_sum(p1, ind1[:, 0], nbatch)
            N = tf.math.unsorted_segment_sum(tf.ones_like(ind1, tf.float32), ind1, tf.reduce_max(ind1)+1)

            p_charge = q_molecule/N
            charge_corr = tf.gather(p_charge, ind1)[:,0]
            p1 =  p1 - charge_corr

        if model_params['neutral_unit'] == 'water_molecule':
            q_molecule = tf.math.reduce_sum(tf.reshape(p1,[-1,3]),axis=1)

            p_charge = q_molecule/3
            charge_corr = tf.reshape(tf.stack([p_charge, p_charge, p_charge], axis=1), [1,-1])[0,:]
            p1 =  p1 - charge_corr

    p1 = tf.expand_dims(p1, axis=1)

    q_tot = tf.math.unsorted_segment_sum(p1, ind1[:, 0], nbatch)

    q_d_a = p1 * features['coord']
    q_d = tf.math.unsorted_segment_sum(q_d_a, ind1[:, 0], nbatch)

    # Compute bond vector
    disp_r = features['diff']

    # Compute atomic dipole
    atomic_d_pairwise = ipred * disp_r
    atomic_d_a = tf.math.unsorted_segment_sum(atomic_d_pairwise, ind2[:, 0], natoms) 
    atomic_d = tf.math.unsorted_segment_sum(atomic_d_a, ind1[:, 0], nbatch)

    a_dipole = q_d_a + atomic_d_a
    dipole = q_d + atomic_d

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, q_tot, ipred, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, q_tot, ipred, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole,
            #'charge': q_tot
            #'atomic_d': tf.expand_dims(a_dipole, 0)
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

pinn.models.AD_BC_R.AD_BC_R_dipole_model

The AD+BC(R) constructs the dipole moment from atomic dipole predictions and atomic pairwise interactions:

\[ \begin{aligned} \mu = \sum_{i}{}^{3}\mathbb{P}_{i} + \sum_{ij}{}^{1}\mathbb{I}_{ij} \cdot \mathbf{r}_{ij} \end{aligned} \]

L2-regularization is applied to the atomic pairwise interactions.

Source code in pinn/models/AD_BC_R.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
@export_model
def AD_BC_R_dipole_model(features, labels, mode, params):
    r"""The AD+BC(R) constructs the dipole moment from 
    atomic dipole predictions and atomic pairwise interactions:

    $$
    \begin{aligned}
    \mu = \sum_{i}{}^{3}\mathbb{P}_{i} + \sum_{ij}{}^{1}\mathbb{I}_{ij} \cdot \mathbf{r}_{ij}
    \end{aligned}
    $$

    L2-regularization is applied to the atomic pairwise interactions.
    """
    network = get_network(params['network'])
    model_params = default_params.copy()
    model_params.update(params['model']['params'])

    features = network.preprocess(features)
    ppred, output_dict = network(features)

    p3 = output_dict['p3']
    p3 = tf.squeeze(p3, axis=-1)
    i1 = output_dict['i1']
    i3 = output_dict['i3']

    ipred = tf.einsum("ixr,ixr->ir", i3, i3) + i1

    ind1 = features['ind_1']  # ind_1 => id of molecule for each atom
    ind2 = features['ind_2']

    natoms = tf.reduce_max(tf.shape(ind1))
    nbatch = tf.reduce_max(ind1)+1 

    p3_d = tf.math.unsorted_segment_sum(p3, ind1[:, 0], nbatch)

    # Compute bond vector
    disp_r = features['diff']

    # Compute atomic dipole
    atomic_d_pairwise = ipred * disp_r
    atomic_d_a = tf.math.unsorted_segment_sum(atomic_d_pairwise, ind2[:, 0], natoms) 
    atomic_d = tf.math.unsorted_segment_sum(atomic_d_a, ind1[:, 0], nbatch)

    a_dipole = p3 + atomic_d_a
    dipole = p3_d + atomic_d

    if model_params['vector_dipole'] == False:
        dipole = tf.sqrt(tf.reduce_sum(dipole**2, axis=1)+1e-6)

    if mode == tf.estimator.ModeKeys.TRAIN:
        metrics = make_metrics(features, dipole, ipred, model_params, mode)
        tvars = network.trainable_variables
        train_op = get_train_op(params['optimizer'], metrics, tvars)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        metrics = make_metrics(features, dipole, ipred, model_params, mode)
        return tf.estimator.EstimatorSpec(mode, loss=tf.reduce_sum(metrics.LOSS),
                                          eval_metric_ops=metrics.METRICS)
    else:
        dipole = dipole / model_params['d_scale']
        dipole *= model_params['d_unit']

        predictions = {
            'dipole': dipole,
            #'charge': q_tot
            #'atomic_d': tf.expand_dims(a_dipole, 0)
        }
        return tf.estimator.EstimatorSpec(
            mode, predictions=predictions)

  1. 1 L. Knijff, and C. Zhang, “Machine learning inference of molecular dipole moment in liquid water,” Mach. Learn.: Sci. Technol. 2(3), 03LT03 (2021). ↩

« Previous
Next »
Built with mkdocs and the flux theme.