def __init__(self,
polymer_length: int,
polymer_dimensions: int,
n_recurrences: int,
mlp_sizes: Tuple[int, ...],
mlp_kwargs: Dict[str, Any] = None,
name: str = 'Energy'):
super().__init__(name=name)
if mlp_kwargs is None:
mlp_kwargs = {
'w_init': hk.initializers.VarianceScaling(),
'b_init': hk.initializers.VarianceScaling(0.1),
'activation': jax.nn.softplus
}
self._graph_net = nn.GraphNetEncoder(n_recurrences, mlp_sizes,
mlp_kwargs)
structure_size = polymer_length * polymer_dimensions
self._decoder = hk.nets.MLP(output_sizes=mlp_sizes +
(structure_size + 1, ),
activate_final=False,
name='GlobalDecoder',
**mlp_kwargs)
def __init__(self, n_recurrences, mlp_sizes, mlp_kwargs=None, name='Energy'):
super(EnergyGraphNet, self).__init__(name=name)
if mlp_kwargs is None:
mlp_kwargs = {
'w_init': hk.initializers.VarianceScaling(),
'b_init': hk.initializers.VarianceScaling(0.1),
'activation': jax.nn.softplus
}
self._graph_net = nn.GraphNetEncoder(n_recurrences, mlp_sizes, mlp_kwargs)
self._decoder = hk.nets.MLP(output_sizes=mlp_sizes + (1,),
activate_final=False,
name='GlobalDecoder',
**mlp_kwargs)