ai2_kit.domain.deepmd module

class ai2_kit.domain.deepmd.CllDeepmdContext(path_prefix: str, resource_manager: ai2_kit.core.resource_manager.ResourceManager, config: ai2_kit.domain.deepmd.CllDeepmdContextConfig)

Bases: BaseCllContext

config: CllDeepmdContextConfig

class ai2_kit.domain.deepmd.CllDeepmdContextConfig(*, script_template: BashTemplate, dp_cmd: str = 'dp', concurrency: int = 5)

Bases: BaseModel

concurrency: int

dp_cmd: str

script_template: BashTemplate

class ai2_kit.domain.deepmd.CllDeepmdInput(config: ai2_kit.domain.deepmd.CllDeepmdInputConfig, mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka'], type_map: List[str], sel_type: Optional[List[int]], old_dataset: List[ai2_kit.core.artifact.Artifact], new_dataset: List[ai2_kit.core.artifact.Artifact])

Bases: object

config: CllDeepmdInputConfig

mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka']

new_dataset: List[Artifact]

old_dataset: List[Artifact]

sel_type: Optional[List[int]]

type_map: List[str]

class ai2_kit.domain.deepmd.CllDeepmdInputConfig(*, train_dw: Optional[DwTraningConfig] = None, model_num: int = 4, init_dataset: List[str] = [], input_template: dict = {}, compress_model: bool = False, isolate_outliers: bool = False, outlier_f_cutoff: float = 10.0, outlier_weight: float = 0.003, fixture_models: List[str] = [], group_by_formula: bool = False)

Bases: BaseModel

class DwTraningConfig(*, input_template: dict)

Bases: BaseModel

Options for deep wannier model training

input_template: dict

Deepmd input template. Ref: https://docs.deepmodeling.com/projects/deepmd/en/master/model/dplr.html

compress_model: bool

Whether to compress model after training.

fixture_models: List[str]

Fixture models used to initialize training, support glob pattern. If this is not empty, then the whole training process will be skipped. This feature is useful for debugging, or explore more structures without training. The models should be on the remote executor. The name fixture is used as the concept of fixture in pytest.

group_by_formula: bool

Grouping dataset by formula If this is enabled, then the dataset will be grouped by formula. Otherwise, the dataset will be grouped by ancestor.

Set this to True when you have multiple structures with the same ancestor.

init_dataset: List[str]

Dataset used to initialize training.

input_template: dict

Deepmd input template.

isolate_outliers: bool

If isolate_outliers is enabled, then outlier data will be separated from training data.

model_num: int

Total number of models to train.

outlier_f_cutoff: float

The threshold of force magnitude to determine whether a data is outlier.

outlier_weight: float

The weight of outlier data in training data.

train_dw: Optional[DwTraningConfig]

Options for deep wannier model training, if None, then deep wannier model will not be trained.

class ai2_kit.domain.deepmd.GenericDeepmdOutput(models: List[ai2_kit.core.artifact.Artifact], dataset: List[ai2_kit.core.artifact.Artifact])

Bases: ICllTrainOutput

dataset: List[Artifact]

get_mlp_models() → List[Artifact]

get_training_dataset() → List[Artifact]

models: List[Artifact]

async ai2_kit.domain.deepmd.cll_deepmd(input: CllDeepmdInput, ctx: CllDeepmdContext)

ai2_kit.domain.deepmd.make_deepmd_dataset(dataset_dir: str, outlier_dir: str, raw_data_collection: ~typing.List[~ai2_kit.core.artifact.__ArtifactDict.<locals>.ArtifactDict], isolate_outliers: bool, outlier_f_cutoff: float, type_map: ~typing.List[str], deepmd_input_template: dict, group_by_formula: bool, mode: str, sel_type: ~typing.Optional[~typing.List[int]])

ai2_kit.domain.deepmd.make_deepmd_input(input_template: dict, type_map: List[str], train_systems: List[str], outlier_systems: List[str], validation_systems: List[str], isolate_outliers: bool, outlier_weight: float)

ai2_kit.domain.deepmd.make_deepmd_task_dirs(input_template: dict, model_num: int, type_map: List[str], train_systems: List[str], outlier_systems: List[str], validation_systems: List[str], isolate_outliers: bool, outlier_weight: float, dw_input_template: Optional[dict], base_dir: str)