# SPDX-License-Identifier: LGPL-3.0-or-later
from typing import (
List,
Optional,
Union,
)
import numpy as np
from deepmd.infer.deep_tensor import (
DeepTensor,
)
[docs]class DeepPolar(DeepTensor):
"""Deep polar model.
Parameters
----------
model_file : Path
The name of the frozen model file.
*args : list
Positional arguments.
auto_batch_size : bool or int or AutoBatchSize, default: True
If True, automatic batch size will be used. If int, it will be used
as the initial batch size.
neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
The ASE neighbor list class to produce the neighbor list. If None, the
neighbor list will be built natively in the model.
**kwargs : dict
Keyword arguments.
"""
@property
def output_tensor_name(self) -> str:
return "polar"
[docs]class DeepGlobalPolar(DeepTensor):
@property
def output_tensor_name(self) -> str:
return "global_polar"
[docs] def eval(
self,
coords: np.ndarray,
cells: Optional[np.ndarray],
atom_types: Union[List[int], np.ndarray],
atomic: bool = False,
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
mixed_type: bool = False,
**kwargs: dict,
) -> np.ndarray:
"""Evaluate the model.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types : list[int] or np.ndarray
The atom types
The list should contain natoms ints
atomic
If True (default), return the atomic tensor
Otherwise return the global tensor
fparam
Not used in this model
aparam
Not used in this model
mixed_type
Whether to perform the mixed_type mode.
If True, the input data has the mixed_type format (see doc/model/train_se_atten.md),
in which frames in a system may have different natoms_vec(s), with the same nloc.
Returns
-------
tensor
The returned tensor
If atomic == False then of size nframes x output_dim
else of size nframes x natoms x output_dim
"""
return super().eval(
coords,
cells,
atom_types,
atomic=atomic,
fparam=fparam,
aparam=aparam,
mixed_type=mixed_type,
**kwargs,
)