def __init__(self,
model,
device="cpu",
collect_triples=False,
environment_provider=SimpleEnvironmentProvider(),
energy=None,
forces=None,
stress=None,
energy_units="eV",
forces_units="eV/Angstrom",
stress_units="eV/Angstrom/Angstrom/Angstrom",
**kwargs):
Calculator.__init__(self, **kwargs)
self.model = model
self.model.to(device)
self.atoms_converter = AtomsConverter(
environment_provider=environment_provider,
collect_triples=collect_triples,
device=device,
)
self.model_energy = energy
self.model_forces = forces
self.model_stress = stress
# Convert to ASE internal units (energy=eV, length=A)
self.energy_units = MDUnits.unit2unit(energy_units, "eV")
self.forces_units = MDUnits.unit2unit(forces_units, "eV/Angstrom")
self.stress_units = MDUnits.unit2unit(stress_units, "eV/A/A/A")
def __init__(self,
model,
device="cpu",
collect_triples=False,
environment_provider=SimpleEnvironmentProvider(),
energy=None,
forces=None,
stress=None,
energy_units="eV",
forces_units="eV/Angstrom",
stress_units="eV/Angstrom/Angstrom/Angstrom",
**kwargs):
Calculator.__init__(self, **kwargs)
self.model = model
self.model.to(device)
self.atoms_converter = AtomsConverter(
environment_provider=environment_provider,
collect_triples=collect_triples,
device=device,
)
self.model_energy = energy
self.model_forces = forces
self.model_stress = stress
# Convert to ASE internal units
# MDUnits parses the given energy units and converts them to atomic units as the common denominator.
# These are then converted to ASE units
self.energy_units = MDUnits.parse_mdunit(energy_units) * units.Ha
self.forces_units = MDUnits.parse_mdunit(
forces_units) * units.Ha / units.Bohr
self.stress_units = (MDUnits.parse_mdunit(stress_units) * units.Ha /
units.Bohr**3)
def __init__(self,
model,
device="cpu",
collect_triples=False,
environment_provider=SimpleEnvironmentProvider(),
energy=None,
forces=None,
energy_units="eV",
forces_units="eV/Angstrom",
**kwargs):
Calculator.__init__(self, **kwargs)
self.model = model
self.atoms_converter = AtomsConverter(
environment_provider=environment_provider,
collect_triples=collect_triples,
device=device,
)
self.model_energy = energy
self.model_forces = forces
# Convert to ASE internal units
self.energy_units = MDUnits.parse_mdunit(energy_units) * units.Ha
self.forces_units = MDUnits.parse_mdunit(
forces_units) * units.Ha / units.Bohr
def __init__(
self,
required_properties,
force_handle,
position_conversion=1.0,
force_conversion=1.0,
property_conversion={},
stress_handle=None,
stress_conversion=1.0,
detach=True,
):
self.results = {}
self.force_handle = force_handle
self.stress_handle = stress_handle
self.required_properties = required_properties
# Perform automatic conversion of units
# internal -> calculator
self.position_conversion = MDUnits.internal2unit(position_conversion)
# calculator -> internal
self.force_conversion = MDUnits.unit2internal(force_conversion)
self.stress_conversion = MDUnits.unit2internal(stress_conversion)
self.property_conversion = {
p: MDUnits.unit2internal(property_conversion[p])
for p in property_conversion
}
# Default conversion (1.0) for all units not set above
self._init_default_conversion()
self.detach = detach
def build_calculator(
_log,
required_properties,
force_handle,
position_conversion,
force_conversion,
property_conversion,
calculator,
device,
):
"""
Build the calculator object from the provided settings.
Args:
model (torch.nn.module): the model which is used for property calculation
required_properties (list): list of properties that are calculated by the model
force_handle (str): name of the forces property in the model output
position_conversion (float): conversion factor for positions
force_conversion (float): conversion factor for forces
property_conversion (dict): dictionary with conversion factors for other properties
calculator (src.schnetpack.md.calculator.Calculator): calculator object
Returns:
the calculator object
"""
_log.info(f"Using {calculator}")
position_conversion = MDUnits.parse_mdunit(position_conversion)
force_conversion = MDUnits.parse_mdunit(force_conversion)
if calculator == "schnet_calculator":
model = load_model(device=device)
return SchnetPackCalculator(
model,
required_properties=required_properties,
force_handle=force_handle,
position_conversion=position_conversion,
force_conversion=force_conversion,
property_conversion=property_conversion,
)
else:
raise NotImplementedError
def __init__(
self,
energy_threshold,
acceleration_factor,
energy_handle="energy",
energy_conversion=1.0,
):
self.energy_threshold = energy_threshold
self.acceleration_factor = acceleration_factor
self.energy_handle = energy_handle
# Convert from calculator -> internal
self.energy_conversion = MDUnits.unit2internal(energy_conversion)
def __init__(
self,
energy_threshold,
acceleration_factor,
energy_handle="energy",
energy_conversion=1.0,
):
self.energy_threshold = energy_threshold
self.acceleration_factor = acceleration_factor
self.energy_handle = energy_handle
# TODO: Implement sensible default behavior for energy conversion
self.energy_conversion = MDUnits.parse_mdunit(energy_conversion)
def __init__(
self,
required_properties,
force_handle,
position_conversion=1.0,
force_conversion=1.0,
property_conversion={},
detach=True,
):
self.results = {}
self.force_handle = force_handle
self.required_properties = required_properties
# Perform automatic conversion of units
self.position_conversion = MDUnits.parse_mdunit(position_conversion)
self.force_conversion = MDUnits.parse_mdunit(force_conversion)
self.property_conversion = {
p: MDUnits.parse_mdunit(property_conversion[p])
for p in property_conversion
}
self._init_default_conversion()
self.detach = detach
class GLEMatrixParser:
"""
General parser for GLE thermostat files. Reads from start string until end of file or a given stop string. If the
argument split is specified, the read matrices are split at the given token. Automatically recognizes
used units and converts them to atomic units.
Args:
start (str): Token when to start reading.
stop (str): Token when to stop reading. If None (default) reads until eno of file.
split (str): If the given token is encountered, matrices are split at this point. If None (default), no split is
performed.
"""
# Automatically recognized format and converts to units
# TODO: think about this carefully
unit_conversions = {
"atomic time units^-1": 1.0 / MDUnits.unit2internal("aut"),
"seconds^-1": 1.0 / MDUnits.unit2internal("s"),
"femtoseconds^-1": 1.0 / MDUnits.unit2internal("fs"),
"picoseconds^-1": 1.0 / 1000.0 / MDUnits.unit2internal("fs"),
"eV": MDUnits.unit2internal("eV"),
"atomic energy units": MDUnits.unit2internal("Ha"),
"K": MDUnits.kB,
}
def __init__(self, start, stop=None, split=None):
self.start = start
self.stop = stop
self.split = split
self.read = False
self.units = None
self._matrix = []
self._tmp_matrix = []
def read_line(self, line):
"""
Read and parse a line obtained from an open file object containing GLE parameters.
Args:
line (str): Line of a GLE parameter file.
"""
line = line.strip()
# Filter for empty lines
if line:
# Check if start token is present
if self.start in line:
self.read = True
# Get units used
unit_name = line.split("(")[-1].replace(")", "")
self.units = self.unit_conversions[unit_name]
elif self.read:
if line.startswith("#"):
# Check for stop and split tokens
if self.stop is not None and self.stop in line:
self.read = False
if self.split is not None and self.split in line:
if len(self._tmp_matrix) > 0:
self._matrix.append(self._tmp_matrix)
self._tmp_matrix = []
else:
# Otherwise read and parse line
self._tmp_matrix.append([float(x) for x in line.split()])
@property
def matrix(self):
"""
Property to get parsed matrices converted to numpy arrays using atomic units.
Returns:
numpy.array: Array of the parsed GLE matrix with the shape normal_modes x s+1 x s+1, where normal_modes is 1
except in the case of the PIGLET thermostat and s is the number of degrees of freedom added via
GLE. If no matrix is found, None is returned.
"""
# Write out last buffer
if len(self._tmp_matrix) > 0:
self._matrix.append(self._tmp_matrix)
# Convert to numpy array
_matrix = np.array(self._matrix)
# Perform unit conversion
if self.units is not None:
return _matrix * self.units
else:
return None