UniformPolarizationModel

class torch_sim.models.polarization.UniformPolarizationModel(device=None, dtype=torch.float64, *, compute_forces=True, compute_stress=True, retain_graph=False)

Bases: ModelInterface

Calculates the energy and force contributions from the application of a constant electric field to a polarizable system.

This model is intended to run after an upstream model inside SerialSumModel.

Required state extras:

• external_E_field

• total_polarization

• polarizability

• born_effective_charges when compute_forces is enabled

Parameters:

• device (device | None)

• dtype (dtype)

• compute_forces (bool)

• compute_stress (bool)

• retain_graph (bool)

property compute_stress: bool

Whether the model computes stresses.

property compute_forces: bool

Whether the model computes forces.

property retain_graph: bool

Whether outputs should remain attached to the autograd graph.

forward(state, **kwargs)

Return additive uniform-field corrections for a polarization model.

Parameters:

state (SimState)

Return type:

dict[str, Tensor]