ai2_kit.domain.lammps module

ai2_kit.domain.lammps module#

class ai2_kit.domain.lammps.CllLammpsContext(path_prefix: str, resource_manager: ai2_kit.core.resource_manager.ResourceManager, config: ai2_kit.domain.lammps.CllLammpsContextConfig)[source]#

Bases: BaseCllContext

config: CllLammpsContextConfig#
class ai2_kit.domain.lammps.CllLammpsContextConfig(*, script_template: BashTemplate, lammps_cmd: str = 'lmp', concurrency: int = 5, multi_gpus_per_job: bool = False)[source]#

Bases: BaseModel

concurrency: int#
lammps_cmd: str#
model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}#

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[dict[str, FieldInfo]] = {'concurrency': FieldInfo(annotation=int, required=False, default=5), 'lammps_cmd': FieldInfo(annotation=str, required=False, default='lmp'), 'multi_gpus_per_job': FieldInfo(annotation=bool, required=False, default=False), 'script_template': FieldInfo(annotation=BashTemplate, required=True)}#

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

multi_gpus_per_job: bool#
script_template: BashTemplate#
class ai2_kit.domain.lammps.CllLammpsInput(config: ai2_kit.domain.lammps.CllLammpsInputConfig, type_map: List[str], mass_map: List[float], mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka'], preset_template: str, new_system_files: List[ai2_kit.core.artifact.Artifact], dp_models: Mapping[str, List[ai2_kit.core.artifact.Artifact]], dp_modifier: Optional[dict], dp_sel_type: Optional[List[int]])[source]#

Bases: object

config: CllLammpsInputConfig#
dp_models: Mapping[str, List[Artifact]]#
dp_modifier: Optional[dict]#
dp_sel_type: Optional[List[int]]#
mass_map: List[float]#
mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka']#
new_system_files: List[Artifact]#
preset_template: str#
type_map: List[str]#
class ai2_kit.domain.lammps.CllLammpsInputConfig(*, n_wise: int = 0, explore_vars: Mapping[str, List[Any]], broadcast_vars: Mapping[str, Any] = {}, preset_template: Optional[str] = None, input_template: Optional[str] = None, template_vars: Mapping[str, Any] = {}, plumed_config: Optional[str] = None, system_files: List[str], ensemble: Optional[Literal['nvt', 'npt', 'npt-i', 'npt-a', 'npt-iso', 'npt-aniso', 'npt-t', 'npt-tri', 'npt-x', 'npt-y', 'npt-z', 'nve', 'csvr']] = None, fix_statement: Optional[str] = None, no_pbc: bool = False, nsteps: int, timestep: float = 0.0005, sample_freq: int = 100, ignore_error: bool = False, type_alias: Mapping[str, List[str]] = {}, custom_ff: Optional[str] = None, fep_opts: FepOptions = FepOptions(ini_ghost_types=[], fin_ghost_types=[]))[source]#

Bases: BaseModel

assert_var(var: str, msg: str = '')[source]#
broadcast_vars: Mapping[str, Any]#

Variants to be explore by broadcast.

Variables defined here won’t join the combination, but will be broadcasted to all combinations.

This can be used to avoid combination explosion.

`yaml LAMBDA_f: [0.0, 0.25, 0.5, 0.75, 1.0] `

custom_ff: Optional[str]#

Allow user to set custom force field. If set, the custom force field will be used instead of the default one. The use can use $$DP_MODELS to reference the deepmd models, and $$SPECORDER to reference the atom type order. For example:

pair_style hybrid/overlay &

deepmd $$DP_MODELS out_freq ${THERMO_FREQ} out_file model_devi.out $$FEP_DP_OPT & buck/coul/long 10.0 10.0

pair_coeff * * deepmd 1 $$SPECORDER pair_coeff * * buck/coul/long 10.0 10.0

ensemble: Optional[Literal['nvt', 'npt', 'npt-i', 'npt-a', 'npt-iso', 'npt-aniso', 'npt-t', 'npt-tri', 'npt-x', 'npt-y', 'npt-z', 'nve', 'csvr']]#
explore_vars: Mapping[str, List[Any]]#

Variants to be explore by full combination or n_wise combination.

Variables defined here will become LAMMPS variables. If multiple value has been set for a variable, the cartesian product will be used to generate the combination. For example,

`yaml TEMP: [330, 430, 530]  # Can be a scalar, e.g. 330 PRES: 1                # Can be a vector, e.g. [1, 2, 3] ` Then you can reference them in the LAMMPS input template as ${TEMP}, ${LAMBDA_f}, ${N_STEPS}, etc.

fep_opts: FepOptions#
fix_statement: Optional[str]#
ignore_error: bool#
input_template: Optional[str]#

LAMMPS input template file content. If set, the preset_template will be ignored.

model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}#

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[dict[str, FieldInfo]] = {'broadcast_vars': FieldInfo(annotation=Mapping[str, Any], required=False, default={}), 'custom_ff': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'ensemble': FieldInfo(annotation=Union[Literal['nvt', 'npt', 'npt-i', 'npt-a', 'npt-iso', 'npt-aniso', 'npt-t', 'npt-tri', 'npt-x', 'npt-y', 'npt-z', 'nve', 'csvr'], NoneType], required=False, default=None), 'explore_vars': FieldInfo(annotation=Mapping[str, List[Any]], required=True), 'fep_opts': FieldInfo(annotation=FepOptions, required=False, default=FepOptions(ini_ghost_types=[], fin_ghost_types=[])), 'fix_statement': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'ignore_error': FieldInfo(annotation=bool, required=False, default=False), 'input_template': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'n_wise': FieldInfo(annotation=int, required=False, default=0), 'no_pbc': FieldInfo(annotation=bool, required=False, default=False), 'nsteps': FieldInfo(annotation=int, required=True), 'plumed_config': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'preset_template': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'sample_freq': FieldInfo(annotation=int, required=False, default=100), 'system_files': FieldInfo(annotation=List[str], required=True), 'template_vars': FieldInfo(annotation=Mapping[str, Any], required=False, default={}), 'timestep': FieldInfo(annotation=float, required=False, default=0.0005), 'type_alias': FieldInfo(annotation=Mapping[str, List[str]], required=False, default={})}#

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

n_wise: int#

The way of combining variants. if n_wise is less than 2 or greater than total fields, the full combination will be used. Or else, the n_wise combination will be used. It is strongly recommended to use n_wise when the full combination is too large.

no_pbc: bool#
nsteps: int#
plumed_config: Optional[str]#

Plumed config file content.

preset_template: Optional[str]#

Name of the preset template.

sample_freq: int#
system_files: List[str]#

Artifacts key of lammps input data

template_vars: Mapping[str, Any]#

input_template may provide extra injection points for user to inject custom settings. Those value could be set here.

Those vars can be referenced in the LAMMPS input template as $$VAR_NAME.

timestep: float#
type_alias: Mapping[str, List[str]]#

Type alias for atoms. For example, if you want to distinguish ghost H and H of HF molecule from other H atoms, you can define the alias as follows: ```yaml type_alias:

H: [ H_ghost, H_hf ]

` And then you can reference them in the LAMMPS input template, for example ` set atom 1 type ${H_hf} set atom 2 type ${H_ghost} ```

classmethod validate_domain(values)[source]#
classmethod validate_explore_variants(value)[source]#
class ai2_kit.domain.lammps.FepOptions(*, ini_ghost_types: List[str] = [], fin_ghost_types: List[str] = [])[source]#

Bases: BaseModel

fin_ghost_types: List[str]#
ini_ghost_types: List[str]#
model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}#

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[dict[str, FieldInfo]] = {'fin_ghost_types': FieldInfo(annotation=List[str], required=False, default=[]), 'ini_ghost_types': FieldInfo(annotation=List[str], required=False, default=[])}#

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

class ai2_kit.domain.lammps.GenericLammpsOutput(model_devi_outputs: List[ai2_kit.core.artifact.Artifact])[source]#

Bases: ICllExploreOutput

get_model_devi_dataset() List[Artifact][source]#
model_devi_outputs: List[Artifact]#
class ai2_kit.domain.lammps.LammpsInputTemplate(template)[source]#

Bases: Template

delimiter = '$$'#
pattern = re.compile('\n            \\$\\$(?:\n              (?P<escaped>\\$\\$)  |   # Escape sequence of two delimiters\n              (?P<named>(?a:[_a-z][_a-z0-9]*))       |   # delimiter and a Python identifier\n    , re.IGNORECASE|re.VERBOSE)#
async ai2_kit.domain.lammps.cll_lammps(input: CllLammpsInput, ctx: CllLammpsContext)[source]#
ai2_kit.domain.lammps.get_ensemble(ensemble: str, group='all')[source]#
ai2_kit.domain.lammps.get_types_template_vars(type_map: List[str], mass_map: List[float], type_alias: Mapping[str, List[str]], sel_type: Optional[List[int]], fep_ini_ghost_types: List[str], fep_fin_ghost_types: List[str])[source]#

generate template vars that related to type_map, mass_map, type_alias, sel_type

the order of atom type index is: real atoms (defined in type_map), virtual atoms (defined in sel_type) and then alias (defined in type_alias)

ai2_kit.domain.lammps.make_lammps_task_dirs(combination_vars: Mapping[str, Sequence[Any]], broadcast_vars: Mapping[str, Sequence[Any]], data_files: List[ArtifactDict], dp_models: Mapping[str, List[str]], n_steps: int, timestep: float, sample_freq: float, no_pbc: bool, n_wise: int, ensemble: Optional[str], fix_statement: Optional[str], preset_template: str, input_template: Optional[str], plumed_config: Optional[str], extra_template_vars: Mapping[str, Any], type_map: List[str], mass_map: List[float], type_alias: Mapping[str, List[str]], work_dir: str, dp_modifier: Optional[dict], dp_sel_type: Optional[List[int]], mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka'], fep_opts: FepOptions, rel_path: bool = False, custom_ff: Optional[str] = None)[source]#