def __init__(
self,
calculator,
nprocs: int = 1,
residual_fn: Optional[Callable] = None,
residual_data: Optional[Dict[str, Any]] = None,
):
if not torch_avail:
report_import_error("pytorch")
default_residual_data = {
"energy_weight": 1.0,
"forces_weight": 1.0,
"stress_weight": 1.0,
"normalize_by_natoms": True,
}
residual_data = _check_residual_data(residual_data, default_residual_data)
_check_compute_flag(calculator, residual_data)
self.calculator = calculator
self.nprocs = nprocs
self.residual_fn = (
energy_forces_residual if residual_fn is None else residual_fn
)
self.residual_data = residual_data
self.optimizer = None
self.optimizer_state_path = None
logger.debug(f"`{self.__class__.__name__}` instantiated.")
def __init__(
self,
model_name: str,
params_relation_callback: Optional[Callable] = None,
):
if not kimpy_avail:
report_import_error("kimpy", self.__class__.__name__)
self.kim_model = self._create_kim_model(model_name)
super(KIMModel, self).__init__(model_name, params_relation_callback)
def __init__(
self,
kim_ca,
config: Configuration,
supported_species: Dict[str, int],
influence_distance: float,
compute_energy: bool = True,
compute_forces: bool = True,
compute_stress: bool = False,
):
if not kimpy_avail:
report_import_error("kimpy", self.__class__.__name__)
# kim compute argument
self.kim_ca = kim_ca
# get supported property
self._get_implemented_property()
super(KIMComputeArguments, self).__init__(
config,
supported_species,
influence_distance,
compute_energy,
compute_forces,
compute_stress,
)
# neighbor list
self.neigh = None
self.num_contributing_particles = None
self.num_padding_particles = None
self.padding_image_of = None
# model input
self.num_particles = None
self.species_code = None
self.particle_contributing = None
self.coords = None
# model output
self.energy = None
self.forces = None
self._init_neigh()
self._update_neigh(influence_distance)
self._register_data(compute_energy, compute_forces)
def _scipy_optimize(self, method, **kwargs):
"""
Minimize the loss use scipy.optimize.least_squares or scipy.optimize.minimize
methods. A user should not call this function, but should call the ``minimize``
method.
"""
size = parallel.get_MPI_world_size()
if size > 1:
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
logger.info(f"Running in MPI mode with {size} processes.")
if self.nprocs > 1:
logger.warning(
f"Argument `nprocs = {self.nprocs}` provided at initialization is "
f"ignored. When running in MPI mode, the number of processes "
f"provided along with the `mpiexec` (or `mpirun`) command is used."
)
x = self.calculator.get_opt_params()
if method in self.scipy_least_squares_methods:
# geodesic LM
if method == "geodesiclm":
if not geodesicLM_avail:
report_import_error("geodesiclm")
else:
minimize_fn = geodesiclm
else:
minimize_fn = scipy.optimize.least_squares
func = self._get_residual_MPI
elif method in self.scipy_minimize_methods:
minimize_fn = scipy.optimize.minimize
func = self._get_loss_MPI
if rank == 0:
result = minimize_fn(func, x, method=method, **kwargs)
# notify other process to break func
break_flag = True
for i in range(1, size):
comm.send(break_flag, dest=i, tag=i)
else:
func(x)
result = None
result = comm.bcast(result, root=0)
return result
else:
# 1. running MPI with 1 process
# 2. running without MPI at all
# both cases are regarded as running without MPI
if self.nprocs == 1:
logger.info("Running in serial mode.")
else:
logger.info(
f"Running in multiprocessing mode with {self.nprocs} processes."
)
# Maybe one thinks he is using MPI because nprocs is used
if mpi4py_avail:
logger.warning(
"`mpi4y` detected. If you try to run in MPI mode, you should "
"execute your code via `mpiexec` (or `mpirun`). If not, ignore "
"this message."
)
x = self.calculator.get_opt_params()
if method in self.scipy_least_squares_methods:
if method == "geodesiclm":
if not geodesicLM_avail:
report_import_error("geodesiclm")
else:
minimize_fn = geodesiclm
else:
minimize_fn = scipy.optimize.least_squares
func = self._get_residual
elif method in self.scipy_minimize_methods:
minimize_fn = scipy.optimize.minimize
func = self._get_loss
result = minimize_fn(func, x, method=method, **kwargs)
return result
def _scipy_optimize(self, method, **kwargs):
"""
Minimize the loss use scipy.optimize.least_squares or scipy.optimize.minimize
methods. A user should not call this function, but should call the ``minimize``
method.
"""
size = parallel.get_MPI_world_size()
if size > 1:
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
msg = "Running in MPI mode with {} processes.".format(size)
log_entry(logger, msg, level="info", print_end="\n\n")
if self.nprocs > 1:
msg = (
'Argument "nprocs = {}" provided at initialization is ignored. When '
"running in MPI mode, the number of processes provided along with "
'the "mpiexec" (or "mpirun") command is used.'.format(
self.nprocs))
log_entry(logger, msg, level="warning")
x = self.calculator.get_opt_params()
if method in self.scipy_least_squares_methods:
# geodesic LM
if method == "geodesiclm":
if not geodesicLM_avail:
report_import_error("geodesciLM")
else:
minimize_fn = geodesiclm
else:
minimize_fn = scipy.optimize.least_squares
func = self._get_residual_MPI
elif method in self.scipy_minimize_methods:
minimize_fn = scipy.optimize.minimize
func = self._get_loss_MPI
if rank == 0:
result = minimize_fn(func, x, method=method, **kwargs)
# notify other process to break func
break_flag = True
for i in range(1, size):
comm.send(break_flag, dest=i, tag=i)
else:
func(x)
result = None
result = comm.bcast(result, root=0)
return result
else:
# 1. running MPI with 1 process
# 2. running without MPI at all
# both cases are regarded as running without MPI
if self.nprocs == 1:
msg = "Running in serial mode."
log_entry(logger, msg, level="info", print_end="\n\n")
else:
msg = "Running in multiprocessing mode with {} processes.".format(
self.nprocs)
log_entry(logger, msg, level="info", print_end="\n\n")
# Maybe one thinks he is using MPI because nprocs is used
if mpi4py_avail:
msg = (
'"mpi4y" detected. If you try to run in MPI mode, you should '
'execute your code via "mpiexec" (or "mpirun"). If not, ignore '
"this message.")
log_entry(logger, msg, level="warning")
x = self.calculator.get_opt_params()
if method in self.scipy_least_squares_methods:
if method == "geodesiclm":
from geodesicLM import geodesiclm
minimize_fn = geodesiclm
else:
minimize_fn = scipy.optimize.least_squares
func = self._get_residual
elif method in self.scipy_minimize_methods:
minimize_fn = scipy.optimize.minimize
func = self._get_loss
result = minimize_fn(func, x, method=method, **kwargs)
return result