xnn_step package#

Submodules#

xnn_step.checkpoint module#

Read the training configuration stored in an xnn checkpoint, without PyTorch.

An xnn checkpoint is what Trainer.save writes with torch.save: {"model": state_dict, "cfg": Config}. On disk that is a zip archive whose <archive>/data.pkl is a pickle of the dictionary, with every tensor stored separately and referenced by a persistent id. The configuration is ordinary Python data (dataclasses holding numbers, strings, lists and dicts), so it can be recovered with a restricted unpickler:

every class the pickle names (xnn...Config, torch.dtype, …) is replaced by an inert stand-in that just records its state, so nothing from the file is executed – unlike torch.load(weights_only=False);
tensor storages are never read (persistent ids resolve to None).

This lets the SEAMM side, which has no PyTorch, see what a model is: its family, cutoff and elements, and how dispersion enters it.

Dispersion conventions (see dispersion_settings()):

cfg.model.extra["dispersion"]: The model was trained with the D3/D4 wrapper, so the checkpoint carries the term and the xnn mdi engine adds it by itself (route A).
cfg.subtracted_dispersion: The model was trained without the wrapper, on targets from which this dispersion term was subtracted (route B). The record is top-level in the configuration, and xnn mdi (xnns 0.4.0 and later) adds the identical term back by itself, so here too the plug-in passes nothing. The value is a mapping of the term’s options with name set, a bare name such as "d4", or True for the D4 defaults.

xnn_step.checkpoint.IN_MODEL_KEY = 'dispersion'#: The cfg.model.extra key of a dispersion term carried by the model (route A).

xnn_step.checkpoint.SUBTRACTED_KEY = 'subtracted_dispersion'#: The top-level cfg key recording the dispersion term subtracted from the training targets, which the engine adds back (route B).

xnn_step.checkpoint.describe_dispersion(spec)[source]#: A short label for a dispersion mapping, e.g. D4 or D4, 12 Å.

xnn_step.checkpoint.dispersion_settings(cfg)[source]#

How dispersion enters a model, from its stored configuration.

Parameters:: cfg (dict or None) – As returned by read_checkpoint_config().
Returns:: (in_model, subtracted): the dispersion term the checkpoint carries, and the term subtracted from the training targets. The engine adds either by itself. True is normalized to {"name": "d4"}, the PBE0-D4 defaults, and a bare name "d3"/"d4" to {"name": ...}.
Return type:: (dict or None, dict or None)
Raises:: ValueError – If the configuration has both: the model was then trained with the wrapper on dispersion-subtracted targets, which removes the dispersion twice, and no run-time setting can correct it.

xnn_step.checkpoint.read_checkpoint_config(path)[source]#

The training configuration stored in an xnn checkpoint, as plain data.

Parameters:: path (str or pathlib.Path) – The checkpoint (.pt) file.
Returns:: The Config as nested dicts (cfg["model"]["extra"] etc.), or None if the file cannot be read as a trainer checkpoint (the reason is logged). Results are cached per file path, size and modification time.
Return type:: dict or None

xnn_step.installer module#

Installer for the xnn plug-in.

This handles any further installation needed after installing the Python package xnn-step: the xnn machine-learned force-field code lives in its own conda environment (seamm-xnn, with PyTorch), which this installer creates from the bundled seamm-xnn.yml and registers in xnn.ini.

class xnn_step.installer.Installer(logger=<Logger xnn_step.installer (WARNING)>)[source]#

Bases: InstallerBase

Handle further installation needed after installing xnn-step.

The Python package xnn-step should already be installed, using pip, conda, or similar. This plug-in-specific installer then checks for the xnn command and its location in xnn.ini:

If it is already registered in xnn.ini there is nothing else to do.
If it can be found using which its location is added to xnn.ini.
Otherwise it is installed in a separate conda environment, seamm-xnn by default, using the seamm-xnn.yml environment file shipped in the plug-in’s data/ directory.

PyTorch is not in that file. The installer puts it in itself, from the PyTorch index that matches the machine’s NVIDIA driver (or PyPI on macOS, for MPS), leaves a torch that works alone, and checks afterwards that torch sees the GPU and that xnn and mdi import. torch-build in xnn.ini (or –torch-tag) forces a build, e.g. on a cluster login node without a GPU. See seamm_manager.torch_support.

exe_version(config)[source]#

Get the version of the xnn package in its environment.

Parameters:: config (dict) – Configuration data for invoking xnn. Typical keys are conda (the path to the conda executable) and conda-environment (the name or full path of the environment in which xnn is installed).
Returns:: The code label and its version, “unknown” if it cannot be found.
Return type:: (“xnn”, str)

xnn_step.metadata module#

This file contains metadata describing the results from Xnn

xnn_step.metadata.metadata = {}#

Description of the computational models for Xnn.

Hamiltonians, approximations, and basis set or parameterizations, only if appropriate for this code. For example:

metadata["computational models"] = {
    "Hartree-Fock": {
        "models": {
            "PM7": {
                "parameterizations": {
                    "PM7": {
                        "elements": "1-60,62-83",
                        "periodic": True,
                        "reactions": True,
                        "optimization": True,
                        "code": "mopac",
                    },
                    "PM7-TS": {
                        "elements": "1-60,62-83",
                        "periodic": True,
                        "reactions": True,
                        "optimization": False,
                        "code": "mopac",
                    },
                },
            },
        },
    },
}

xnn_step.tk_xnn module#

The graphical part of a xnn step

class xnn_step.tk_xnn.TkXnn(tk_flowchart=None, node=None, canvas=None, x=None, y=None, w=200, h=50)[source]#

Bases: TkNode

The graphical part of a xnn step in a flowchart.

tk_flowchart#

The flowchart that we belong to.

Type:: TkFlowchart = None

node#

The corresponding node of the non-graphical flowchart

Type:: Node = None

namespace#

The namespace of the current step.

Type:: str

tk_subflowchart#

A graphical Flowchart representing a subflowchart

Type:: TkFlowchart

canvas#

The Tk Canvas to draw on

Type:: tkCanvas = None

dialog#

The Pmw dialog object

Type:: Dialog

x#

The x-coordinate of the center of the picture of the node

Type:: int = None

y#

The y-coordinate of the center of the picture of the node

Type:: int = None

w#

The width in pixels of the picture of the node

Type:: int = 200

h#

The height in pixels of the picture of the node

Type:: int = 50

self[widget]#

A dictionary of tk widgets built using the information contained in xnn_parameters.py

Type:: dict

See also

Xnn, TkXnn, XnnParameters

create_dialog()[source]#

Create the dialog. A set of widgets will be chosen by default based on what is specified in the xnn_parameters module.

Parameters:: None
Return type:: None

See also

TkXnn.reset_dialog

reset_dialog(widget=None)[source]#

Layout the widgets in the dialog.

The widgets are chosen by default from the information in xnn_parameter.

This function simply lays them out row by row with aligned labels. You may wish a more complicated layout that is controlled by values of some of the control parameters. If so, edit or override this method

Parameters:: widget (Tk Widget = None)
Return type:: None

See also

TkXnn.create_dialog

right_click(event)[source]#

Handles the right click event on the node.

Parameters:: event (Tk Event)
Return type:: None

See also

TkXnn.edit

xnn_step.xnn module#

Non-graphical part of the xnn step in a SEAMM flowchart.

The xnn plug-in provides machine-learned force fields (MLFFs) trained with xnn as model chemistries: the Model Chemistry step lists the checkpoints found in the directories configured in xnn.ini (as xnn:MLFF@<model>), and steps that drive a model chemistry over MDI – Energy, LAMMPS, Dimer Builder, Normal Mode Sampling – launch xnn mdi through xnn_step.XnnStep.

There is therefore no xnn step to place in a flowchart. This node exists only to satisfy the plug-in interface; running it explains how to use the models.

class xnn_step.xnn.Xnn(flowchart=None, title='xnn', extension=None, logger=<Logger xnn_step.xnn (WARNING)>)[source]#

Bases: Node

The non-graphical part of the xnn step in a flowchart.

See also

TkXnn, XnnParameters, XnnStep

analyze(indent='', **kwargs)[source]#: Nothing to analyze: this step produces no results.

description_text(P=None)[source]#: Create the text description of what this step will do.

property git_revision#: The git version of this module.

run()[source]#: Run the xnn step – which only explains how the plug-in is used.

property version#: The semantic version of this module.

xnn_step.xnn_parameters module#

Control parameters for the xnn step in a SEAMM flowchart

class xnn_step.xnn_parameters.XnnParameters(defaults={}, data=None)[source]#

Bases: Parameters

The control parameters for the xnn step.

The xnn plug-in is a provider of model chemistries (its MLFF checkpoints) for the Model Chemistry step, and the launcher of the xnn mdi engine for the steps that drive one (Energy, LAMMPS). It has no user-facing step of its own, and so no control parameters beyond the standard results.

parameters = {'results': {'default': {}, 'default_units': None, 'description': 'results', 'enumeration': (), 'format_string': '', 'help_text': 'The results to save to variables or in tables.', 'kind': 'dictionary'}}#

xnn_step.xnn_step module#

Helper class needed for the stevedore integration, plus the Model Chemistry / MDI engine contract for xnn machine-learned force fields (MLFFs).

The xnn plug-in is a provider: it advertises the trained xnn checkpoints found in the configured model directories to the Model Chemistry step (as xnn:MLFF@<model>), and builds the command that launches xnn mdi – the checkpoint served as a resident MDI engine – for any step that drives a model chemistry over MDI (Energy, LAMMPS, Dimer Builder, Normal Mode Sampling, …).

A committee – a directory of members m0/best.pt, m1/best.pt, …, as the Active Learning step’s Train Committee writes – is offered the same way, as one model. It runs as the committee-mean engine of the active-learning-step package, which serves the mean of the members and writes how much they disagree.

class xnn_step.xnn_step.XnnStep(flowchart=None, gui=None)[source]#

Bases: object

Helper class needed for the stevedore integration.

This must provide a description() method that returns a dict containing a description of this node, and create_node() and create_tk_node() methods for creating the graphical and non-graphical nodes.

The dictionary for the description is the class variable just below these comments. The felds are as follows:

my_description{str, str}
A human-readable description of this step. It can be several lines long, and needs to be clear to non-expert users. It contains the following keys: description, group, name.

my_description[“description”]tuple
A description of the xnn step. It must be clear to non-experts.

my_description[“group”]str
Which group in the menus to put this step. If the group does not exist it will be created. Common groups are “Building”, “Control”, “Custom”, “Data”, and “Simulations”.

my_description[“name”]str
The name of this step, to be displayed in the menus.

classmethod available_models(config=None, with_sources=False)[source]#

The xnn checkpoints found in the configured model directories.

Parameters:

config (dict, optional) – An xnn.ini section (see get_executor_config()). Defaults to the [local] section of the user’s xnn.ini, or the plug-in’s default ini if the user has none.
with_sources (bool) – Also return, per model, the human-readable source such as personal:xnn/water.pt (or the absolute path for a plain directory).

Returns:

Model name -> checkpoint path (a committee’s directory for a committee), sorted by name. The directories are searched in the order listed, so a personal model shadows a local one of the same name; the shadowed file is logged and ignored.

Return type:

dict[str, Path] (or (dict[str, Path], dict[str, str]) with sources)

static committee_members(path)[source]#

The members’ checkpoints if path is a committee, else None.

A committee is a directory of members m0, m1, … each holding a best.pt, with at least two members. A committee.json in it may list the members to use, e.g. {"members": ["m0", "m1", "m3"]} to leave one out.

Returns:: The checkpoints, in the order of the members.
Return type:: [Path] or None

create_node(flowchart=None, **kwargs)[source]#

Create and return the new node object.

Parameters:

flowchart (seamm.Node) – A non-graphical SEAMM node
**kwargs (keyworded arguments) – Various keyworded arguments such as title, namespace or extension representing the title displayed in the flowchart, the namespace for the plugins of a subflowchart and the extension, respectively.

Return type:

Xnn

create_tk_node(canvas=None, **kwargs)[source]#

Create and return the graphical Tk node object.

Parameters:

canvas (tk.Canvas) – The Tk Canvas widget
**kwargs (keyworded arguments) – Various keyworded arguments such as tk_flowchart, node, canvas, x, y, w, h, etc.

Return type:

TkXnn

description()[source]#: Return a description of what this extension does.

classmethod get_executor_config(executor, seamm_options)[source]#

Return how to launch the xnn MDI engine on this machine.

Reads the per-plug-in xnn.ini for the current executor type, so the engine runs in the conda environment holding xnn and PyTorch (seamm-xnn by default).

Parameters:

executor (seamm.ExecutorBase) – The flowchart executor (self.flowchart.executor in the driver); executor.name selects the ini section.
seamm_options (dict) – The global SEAMM options (self.global_options); seamm_options["root"] locates the ini.

Returns:

The ini section for the current executor, plus version (the plug-in version).

Return type:

dict

classmethod get_mdi_engine_command(executor, seamm_options, *, method, port, hostname='localhost', charge=0, multiplicity=1, n_atoms=None, engine_name='XNN', extra_args=None, mdi_method='TCP', checkpoint=None, device=None, **kwargs)[source]#

Build the argv that launches the xnn MDI engine, over TCP or MPI.

The driver (e.g. the Energy step, via seamm_mdi.MDIEngine) owns the rendezvous (TCP, port, hostname) and passes it in; everything xnn-specific – the conda environment, the checkpoint path for the model, the torch device – is supplied here, so the driver hardwires no xnn knowledge. The atom count, atomic numbers, coordinates and cell all arrive over the MDI handshake from the driver.

Parameters:

executor – Passed straight to get_executor_config.
seamm_options – Passed straight to get_executor_config.
method (str) – The model name, as advertised by get_model_chemistry_options.
port (int) – TCP port the engine dials; chosen by the driver.
hostname (str) – Host the engine dials.
charge (int or float) – The total charge of the system. When nonzero it is passed as --total-charge, which the D4 EEQ charges of a dispersion model need (models without dispersion ignore it). Requires an xnn whose mdi command has that option.
multiplicity (int) – Accepted for interface compatibility; an MLFF has no notion of it.
n_atoms (int, optional) – Number of atoms (accepted for interface compatibility).
engine_name (str) – The MDI -name for the engine (default “XNN”).
extra_args (list of str, optional) – Extra xnn mdi flags appended verbatim (e.g. --dtype float32, --eeq-reuse): the engine options of the run.
mdi_method (str) – "TCP" (the default) for an engine the driver reaches over a port, or "MPI" for one launched with the driver as an MPI pair (LAMMPS’s MLFF runs). With MPI the engine runs in the driver’s environment – the two must share MPI and the MDI library – so the argv has no conda run and no thread-count prefix, and port and hostname are not used.
checkpoint (str or Path, optional) – The checkpoint to use instead of looking method up in the model directories, e.g. one given by path in the Model Chemistry step. A committee’s directory runs the committee-mean engine of the active-learning-step package instead of xnn mdi: it serves the members’ mean and writes their spreads to committee_spreads.csv.
device (str, optional) – The torch device (e.g. cuda:0); the device in xnn.ini, else cpu, by default.

Notes

The dispersion settings come from the checkpoint’s stored training configuration (see xnn_step.checkpoint). The engine adds the dispersion by itself in both cases: a model trained with the dispersion wrapper carries the term, and a model trained on dispersion-subtracted targets records it as cfg.subtracted_dispersion, which xnn mdi adds back. A model with both removes the dispersion twice and is refused.

Returns:: A ready-to-run argv, possibly with a leading OMP_NUM_THREADS=n environment prefix (honored by seamm_mdi.MDIEngine and by shlex.join into a launch script).
Return type:: list of str

classmethod get_model_chemistry_options(periodic_only=False, mdi_only=False)[source]#

Return the model chemistries (level specs) the xnn models provide.

Advertises bare xnn:MLFF@<model> level specs, one per checkpoint found in the model directories configured in xnn.ini, keyed by the model name (the checkpoint’s file stem, with any grammar-reserved characters replaced by -).

Parameters:

periodic_only (bool) – Only return models validated for periodic systems (all of them).
mdi_only (bool) – Only return models launchable via MDI (all of them).

Returns:

Keyed by model name. Each entry carries model_chemistry, type ("MLFF"), description, periodic_native, periodic_mdi, elements (the atomic numbers the model was trained on, e.g. "1,8"; "" if unknown), mdi_capable, mdi_method_arg (the model name), path (the checkpoint), family (e.g. "mace") and dispersion ("", or how dispersion enters, e.g. "D4 in the model"), the last three read from the checkpoint’s stored training configuration.

Return type:

dict

Notes

A checkpoint trained with a dispersion wrapper on dispersion-subtracted targets removes the dispersion twice; it is logged and not offered.

classmethod model_chemistry_for_checkpoint(path)[source]#

The model chemistry for a checkpoint given by its path rather than found in the model directories – e.g. one written during training. Its method is the file’s stem, as for the models found. The path may also be a committee’s directory.

Raises FileNotFoundError if there is no such file, and ValueError for a checkpoint that removes the dispersion twice.

classmethod model_chemistry_option(name, path, source=None)[source]#: The model chemistry for one checkpoint, as get_model_chemistry_options offers it. Raises ValueError for a checkpoint that removes the dispersion twice.

classmethod model_directories(config)[source]#

The directories to search for checkpoints, from models in xnn.ini.

Each line is a directory: an absolute or ~ path, {root}/... for the SEAMM root, or personal:<subdir> for ~/.seamm.d/data/Forcefields/<subdir> and local:<subdir> for <root>/data/Forcefields/<subdir> then ~/SEAMM/data/Forcefields/<subdir> (the same convention as the Forcefield step’s personal:/local: forcefield files). Returns [(label_prefix, Path)] in order; the label prefix (e.g. personal:xnn/) is what a model’s file is reported as.

static model_name(path)[source]#

The model-chemistry method name for a checkpoint.

A checkpoint with a name of its own (water_mace.pt) is called by its file stem. One with a generic name – training writes runs/<...>/<run>/best.pt and last.pt – is called after its directory, with the grandparent too when the directory is only an index (m1, fold2, seed0), keeping the stem unless it is best, and then the first 6 hex digits of the SHA-256 of the file: e.g. x013_m1-3f9a2c, committee_w1-m1-last-7b1e09. The name is that of these weights: a retrained or later checkpoint in the same place gets a different one. The grammar’s reserved characters become -.

A committee is called by its directory’s name. When that is generic (committee), it is called after the nearest directory above it that is neither an index nor generic – a job’s step directories (2/1) are indices – with the start of a hash of its members: e.g. Job_012130/2/1/committee is Job_012130-committee-5c20d7.

my_description = {'description': 'Machine-learned force fields (MLFFs) trained with xnn, provided as model chemistries and run as MDI engines.', 'group': 'Simulations', 'name': 'xnn'}#

static seamm_ini_path()[source]#: Path to the global seamm.ini: ~/.seamm.d/seamm.ini, or $SEAMM_INI if set (handy for testing against an alternate config).

static seamm_root()[source]#

The SEAMM root directory (--root), as a Path.

The parsed --root when available (inside a running flowchart or the editor), else the installation this Python belongs to, SEAMM_ROOT or ~/SEAMM – see seamm_util.current_root.

classmethod thread_count()[source]#: The CPU thread cap for the engine from [xnn-step] ncores in seamm.ini, or None for ‘available’ (PyTorch’s default).

Module contents#

xnn_step A SEAMM plug-in for xnn