Skip to content

Python API

Documentation for the core functions in gpuma.

Package layout

Everything documented here is re-exported from the top-level gpuma namespace, so the public API is flat: from gpuma import optimize_single_smiles, Config, read_multi_xyz all work regardless of where a symbol lives internally. Internally the package is organised into subpackages:

  • gpuma.models — model dispatchers (load_calculator, load_torchsim_model) and the AVAILABLE_*_MODELS registries, with backend-specific loaders under models/fairchem.py, models/orb.py, models/sevennet.py.
  • gpuma.conformer_generation — SMILES → 3D structure/ensemble embedding (embed.py, mol_utils.py).
  • gpuma.utils — I/O (io_handler.py), logging (logging_utils.py) and timing helpers (decorators.py).
  • Top-level modules structure.py, config.py, optimizer.py, api.py, cli.py.

The fully-qualified paths below (e.g. gpuma.utils.io_handler.read_xyz) are the canonical source locations; the short top-level aliases (gpuma.read_xyz) are equivalent.

Public high-level Python API for common geometry optimization workflows.

This module provides convenience functions built on top of the lower-level I/O and optimization utilities. It allows users to easily optimize molecular structures starting from SMILES strings or XYZ files, as well as optimizing ensembles of conformers.

Data Structures

The core container passed between I/O, conformer generation and optimization.

Structure dataclass

Container for a molecular structure used in GPUMA.

Attributes

symbols : list[str] List of atomic symbols. coordinates : list[tuple[float, float, float]] N x 3 list of floats for atomic positions in Angstrom. charge : int Total charge of the system. multiplicity : int Spin multiplicity of the system. energy : float | None Optional energy value of the structure in eV. comment : str Optional comment or metadata string. metadata : dict Free-form metadata dictionary for additional information.

Source code in src/gpuma/structure.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
@dataclass
class Structure:
    """Container for a molecular structure used in GPUMA.

    Attributes
    ----------
    symbols : list[str]
        List of atomic symbols.
    coordinates : list[tuple[float, float, float]]
        ``N x 3`` list of floats for atomic positions in Angstrom.
    charge : int
        Total charge of the system.
    multiplicity : int
        Spin multiplicity of the system.
    energy : float | None
        Optional energy value of the structure in eV.
    comment : str
        Optional comment or metadata string.
    metadata : dict
        Free-form metadata dictionary for additional information.
    """

    symbols: list[str]
    coordinates: list[tuple[float, float, float]]
    charge: int
    multiplicity: int
    energy: float | None = None
    comment: str = ""

    # Room for future metadata without breaking the public API
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def n_atoms(self) -> int:
        """Return the number of atoms in the structure.

        Returns
        -------
        int
            Number of atoms.
        """
        return len(self.symbols)

    def with_energy(self, energy: float | None) -> "Structure":
        """Set the energy of the structure and return the modified instance.

        Parameters
        ----------
        energy : float | None
            Energy value in eV to assign to this structure. ``None`` clears
            the current energy.

        Returns
        -------
        Structure
            The same :class:`Structure` instance, to allow method chaining.
        """
        self.energy = energy
        return self

n_atoms property

Return the number of atoms in the structure.

Returns

int Number of atoms.

with_energy(energy)

Set the energy of the structure and return the modified instance.

Parameters

energy : float | None Energy value in eV to assign to this structure. None clears the current energy.

Returns

Structure The same :class:Structure instance, to allow method chaining.

Source code in src/gpuma/structure.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def with_energy(self, energy: float | None) -> "Structure":
    """Set the energy of the structure and return the modified instance.

    Parameters
    ----------
    energy : float | None
        Energy value in eV to assign to this structure. ``None`` clears
        the current energy.

    Returns
    -------
    Structure
        The same :class:`Structure` instance, to allow method chaining.
    """
    self.energy = energy
    return self

Single Structure Optimization

Methods for optimizing individual molecules provided as SMILES or XYZ files.

optimize_single_smiles(smiles, output_file=None, config=None)

Optimize a single molecule from a SMILES string.

This function uses the provided SMILES string to generate an initial 3D structure using the Morfeus library. It then optimizes the structure using the specified optimization pipeline.

Parameters:

Name Type Description Default
smiles str

SMILES string of the molecule to optimize.

 required
output_file str

Path to an output XYZ file where the optimized structure will be written. If None, the optimized structure is not saved to a file.

 None
config Config

Config object to control the optimization pipeline. Highly recommended to specify. If None, the configuration will be loaded from the default file.

 None

Returns:

Name Type Description
Structure Structure

The optimized molecular structure as a Structure object.

Raises:

Type Description
ValueError

If the generated structure is not valid.
Source code in src/gpuma/api.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def optimize_single_smiles(
    smiles: str,
    output_file: str | None = None,
    config: Config | None = None,
) -> Structure:
    """Optimize a single molecule from a SMILES string.

    This function uses the provided SMILES string to generate an initial 3D structure
    using the Morfeus library. It then optimizes the structure using the specified
    optimization pipeline.

    Args:
        smiles (str): SMILES string of the molecule to optimize.
        output_file (str): Path to an output XYZ file where the optimized structure
            will be written. If None, the optimized structure is not saved to a file.
        config (Config, optional): Config object to control the optimization pipeline.
            Highly recommended to specify. If None, the configuration will be loaded
            from the default file.

    Returns:
        Structure: The optimized molecular structure as a Structure object.

    Raises:
        ValueError: If the generated structure is not valid.

    """
    if config is None:
        config = load_config_from_file()

    structure = smiles_to_xyz(
        smiles, multiplicity=config.optimization.multiplicity, config=config
    )

    if not isinstance(structure, Structure):
        raise ValueError("smiles_to_xyz did not return a Structure")
    structure.comment = f"Optimized from SMILES: {smiles}"
    result = optimize_single_structure(structure, config)

    if output_file:
        save_xyz_file(result, output_file)

    return result

optimize_single_xyz_file(input_file, output_file=None, config=None)

Optimize a single structure from an XYZ file.

This function reads a molecular structure from the specified XYZ file, optimizes it using the provided optimization pipeline, and optionally writes the optimized structure to an output XYZ file.

Parameters:

Name Type Description Default
input_file str

Path to an input XYZ file from which to read the initial structure.

 required
output_file str

Path to an output XYZ file where the optimized structure will be written. If None, the optimized structure will not be saved to a file.

 None
config Config

Config object to control the optimization pipeline. Highly recommended to specify. If None, the configuration will be loaded from the default file.

 None

Returns:

Name Type Description
Structure Structure

The optimized molecular structure as a Structure object.

Raises:

Type Description
ValueError

If the input file does not exist or if the read structure is not valid.
Source code in src/gpuma/api.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def optimize_single_xyz_file(
    input_file: str,
    output_file: str | None = None,
    config: Config | None = None,
) -> Structure:
    """Optimize a single structure from an XYZ file.

    This function reads a molecular structure from the specified XYZ file,
    optimizes it using the provided optimization pipeline, and optionally
    writes the optimized structure to an output XYZ file.

    Args:
        input_file (str): Path to an input XYZ file from which to read the initial structure.
        output_file (str): Path to an output XYZ file where the optimized structure will be written.
            If None, the optimized structure will not be saved to a file.
        config (Config, optional): Config object to control the optimization pipeline.
            Highly recommended to specify. If None, the configuration will be loaded from
            the default file.

    Returns:
        Structure: The optimized molecular structure as a Structure object.

    Raises:
        ValueError: If the input file does not exist or if the read structure is not valid.

    """
    if not os.path.isfile(input_file):
        raise ValueError(f"Input file {input_file} does not exist.")
    if config is None:
        config = load_config_from_file()

    structure = read_xyz(
        input_file,
        charge=int(config.optimization.charge),
        multiplicity=int(config.optimization.multiplicity),
    )
    if not isinstance(structure, Structure):
        raise ValueError("read_xyz did not return a Structure")
    structure.comment = f"Optimized from: {input_file}"
    result = optimize_single_structure(structure, config)

    if output_file:
        save_xyz_file(result, output_file)

    return result

Batch & Ensemble Optimization

Methods for processing multiple structures, ensembles, or entire directories.

optimize_ensemble_smiles(smiles, output_file=None, config=None)

Optimize a conformer ensemble generated from a SMILES string.

This function generates a specified number of conformers from the provided SMILES string using the Morfeus library. It then optimizes each conformer using the specified optimization pipeline. Optionally, the optimized ensemble can be saved to a multi-structure XYZ file.

Parameters:

Name Type Description Default
smiles str

SMILES string of the molecule for which to generate conformers.

 required
output_file str

Path to an output multi-structure XYZ file where the optimized ensemble will be written. If None, the ensemble is not saved to a file.

 None
config Config

Config object to control the optimization pipeline. Highly recommended to specify. If None, the configuration will be loaded from the default file.

 None

Returns:

Type Description
list[Structure]

list[Structure]: A list of optimized molecular structures as Structure objects.

Raises:

Type Description
ValueError

If output_file is not specified when required or if the generated conformers are not valid.
Source code in src/gpuma/api.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def optimize_ensemble_smiles(
    smiles: str,
    output_file: str | None = None,
    config: Config | None = None,
) -> list[Structure]:
    """Optimize a conformer ensemble generated from a SMILES string.

    This function generates a specified number of conformers from the provided
    SMILES string using the Morfeus library. It then optimizes each conformer
    using the specified optimization pipeline. Optionally, the optimized ensemble
    can be saved to a multi-structure XYZ file.

    Args:
        smiles (str): SMILES string of the molecule for which to generate conformers.
        output_file (str, optional): Path to an output multi-structure XYZ file where
            the optimized ensemble will be written. If None, the ensemble is not saved to a file.
        config (Config, optional): Config object to control the optimization pipeline.
            Highly recommended to specify. If None, the configuration will be loaded from
            the default file.

    Returns:
        list[Structure]: A list of optimized molecular structures as Structure objects.

    Raises:
        ValueError: If output_file is not specified when required or if the generated
            conformers are not valid.

    """
    if config is None:
        config = load_config_from_file()
    multiplicity = int(config.optimization.multiplicity)
    num_conformers = int(config.conformer_generation.max_num_conformers)
    seed = int(config.conformer_generation.conformer_seed)
    conformers = smiles_to_ensemble(
        smiles, num_conformers, multiplicity, seed=seed, config=config
    )
    if not isinstance(conformers, list) or (
        len(conformers) and not isinstance(conformers[0], Structure)
    ):
        raise ValueError("smiles_to_ensemble did not return a list of Structure")
    for s in conformers:
        s.multiplicity = multiplicity
    results = optimize_structure_batch(conformers, config)

    if output_file:
        comments = [
            f"Optimized conformer {i + 1} from SMILES: {smiles}" for i in range(len(results))
        ]
        save_multi_xyz(results, output_file, comments)

    return results

optimize_batch_multi_xyz_file(input_file, output_file=None, config=None)

Optimize a batch of structures from a multi-structure XYZ file.

This function reads multiple molecular structures from the specified multi-structure XYZ file, optimizes each structure using the provided optimization pipeline, and writes the optimized structures to an output multi-structure XYZ file.

Parameters:

Name Type Description Default
input_file str

Path to an input multi-structure XYZ file from which to read the initial structures.

 required
output_file str

Path to an output multi-structure XYZ file where the optimized structures will be written. If None, the optimized structures will not be saved to a file.

 None
config Config

Config object to control the optimization pipeline. Highly recommended to specify. If None, the configuration will be loaded from the default file.

 None

Returns:

Type Description
list[Structure]

list[Structure]: A list of optimized molecular structures as Structure objects.

Raises:

Type Description
ValueError

If the input file does not exist or if the read structures are not valid.
Source code in src/gpuma/api.py 
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def optimize_batch_multi_xyz_file(
    input_file: str,
    output_file: str | None = None,
    config: Config | None = None,
) -> list[Structure]:
    """Optimize a batch of structures from a multi-structure XYZ file.

    This function reads multiple molecular structures from the specified multi-structure
    XYZ file, optimizes each structure using the provided optimization pipeline, and writes
    the optimized structures to an output multi-structure XYZ file.

    Args:
        input_file (str): Path to an input multi-structure XYZ file from which to
            read the initial structures.
        output_file (str): Path to an output multi-structure XYZ file where the optimized
            structures will be written.
            If None, the optimized structures will not be saved to a file.
        config (Config, optional): Config object to control the optimization pipeline.
            Highly recommended to specify. If None, the configuration will be loaded from
            the default file.

    Returns:
        list[Structure]: A list of optimized molecular structures as Structure objects.

    Raises:
        ValueError: If the input file does not exist or if the read structures are not valid.

    """
    if not os.path.isfile(input_file):
        raise ValueError(f"Input file {input_file} does not exist.")

    if config is None:
        config = load_config_from_file()

    structures = read_multi_xyz(
        input_file,
        charge=int(config.optimization.charge),
        multiplicity=int(config.optimization.multiplicity),
    )
    results = optimize_structure_batch(structures, config)

    if output_file:
        comments = [
            f"Optimized structure {i + 1} from: {input_file}" for i in range(len(results))
        ]
        save_multi_xyz(results, output_file, comments)

    return results

optimize_batch_xyz_directory(input_directory, output_file, config=None)

Optimize a batch of structures from XYZ files in a directory.

This function reads multiple molecular structures from XYZ files in the specified input directory, optimizes each structure using the provided optimization pipeline, and writes the optimized structures to a multi-structure XYZ output file.

Parameters:

Name Type Description Default
input_directory str

Path to an input directory containing XYZ files.

 required
output_file str

Path to an output multi-structure XYZ file where the optimized structures will be written.

 required
config Config

Config object to control the optimization pipeline. Highly recommended to specify. If None, the configuration will be loaded from the default file.

 None

Returns:

Type Description
list[Structure]

list[Structure]: A list of optimized molecular structures as Structure objects.

Raises:

Type Description
ValueError

If the input directory does not exist or contains no valid XYZ files.
Source code in src/gpuma/api.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def optimize_batch_xyz_directory(
    input_directory: str,
    output_file: str,
    config: Config | None = None,
) -> list[Structure]:
    """Optimize a batch of structures from XYZ files in a directory.

    This function reads multiple molecular structures from XYZ files in the specified input
    directory, optimizes each structure using the provided optimization pipeline,
    and writes the optimized structures to a multi-structure XYZ output file.

    Args:
        input_directory (str): Path to an input directory containing XYZ files.
        output_file (str): Path to an output multi-structure XYZ file where the
            optimized structures will be written.
        config (Config, optional): Config object to control the optimization pipeline.
            Highly recommended to specify. If None, the configuration will be loaded from the
            default file.

    Returns:
        list[Structure]: A list of optimized molecular structures as Structure objects.

    Raises:
        ValueError: If the input directory does not exist or contains no valid XYZ files.

    """
    if config is None:
        config = load_config_from_file()

    structures = read_xyz_directory(
        input_directory,
        charge=int(config.optimization.charge),
        multiplicity=int(config.optimization.multiplicity),
    )
    results = optimize_structure_batch(structures, config)

    if output_file:
        comments = [
            f"Optimized structure {i + 1} from batch input"
            for i in range(len(results))
        ]
        save_multi_xyz(results, output_file, comments)

    return results

Model Loading

Functions for directly loading model calculators and torch-sim wrappers.

load_calculator(config)

Load an ASE-compatible calculator for single-structure optimization.

Dispatches to the Fairchem, ORB-v3, or SevenNet backend based on config.model.model_type.

Parameters

config : Config GPUMA configuration object.

Returns

calculator An ASE calculator (FAIRChemCalculator, ORBCalculator, or SevenNetCalculator).

Raises

ImportError If the required backend package is not installed. ValueError If the model name is unknown or missing.

Source code in src/gpuma/models/__init__.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
@time_it
def load_calculator(config: Config):
    """Load an ASE-compatible calculator for single-structure optimization.

    Dispatches to the Fairchem, ORB-v3, or SevenNet backend based on
    ``config.model.model_type``.

    Parameters
    ----------
    config : Config
        GPUMA configuration object.

    Returns
    -------
    calculator
        An ASE calculator (``FAIRChemCalculator``, ``ORBCalculator``, or
        ``SevenNetCalculator``).

    Raises
    ------
    ImportError
        If the required backend package is not installed.
    ValueError
        If the model name is unknown or missing.
    """
    model_type = resolve_model_type(config)
    if model_type == "orb":
        return _load_orb_calculator(config)
    if model_type == "sevennet":
        return _load_sevennet_calculator(config)
    return _load_fairchem_calculator(config)

load_torchsim_model(config)

Load a torch-sim model wrapper for GPU-accelerated batch optimization.

Dispatches to the Fairchem, ORB-v3, or SevenNet backend based on config.model.model_type.

Parameters

config : Config GPUMA configuration object.

Returns

model A torch-sim model (FairChemModel, OrbTorchSimModel, or a Float64Wrapper-wrapped SevenNet model).

Raises

ImportError If the required backend package is not installed. ValueError If the model name is unknown or missing.

Source code in src/gpuma/models/__init__.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
@time_it
def load_torchsim_model(config: Config):
    """Load a torch-sim model wrapper for GPU-accelerated batch optimization.

    Dispatches to the Fairchem, ORB-v3, or SevenNet backend based on
    ``config.model.model_type``.

    Parameters
    ----------
    config : Config
        GPUMA configuration object.

    Returns
    -------
    model
        A torch-sim model (``FairChemModel``, ``OrbTorchSimModel``, or a
        ``Float64Wrapper``-wrapped SevenNet model).

    Raises
    ------
    ImportError
        If the required backend package is not installed.
    ValueError
        If the model name is unknown or missing.
    """
    model_type = resolve_model_type(config)
    if model_type == "orb":
        return _load_orb_torchsim(config)
    if model_type == "sevennet":
        return _load_sevennet_torchsim(config)
    return _load_fairchem_torchsim(config)

Configuration

The Config object controls model selection, optimization settings, conformer generation and technical/device options. It is accepted by every high- and low-level optimization function. A ready-to-use instance carrying the built-in defaults is available as gpuma.default_config (equivalent to Config()). See the Configuration page for the full list of keys and defaults.

Config

Dict-backed configuration with attribute access for sections.

Example:

cfg = load_config_from_file() print(cfg.technical.logging_level) cfg.technical.device = "cuda" save_config_to_file(cfg, "config.json")

Source code in src/gpuma/config.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
class Config:
    """Dict-backed configuration with attribute access for sections.

    Example:
    -------
    >>> cfg = load_config_from_file()
    >>> print(cfg.technical.logging_level)
    >>> cfg.technical.device = "cuda"
    >>> save_config_to_file(cfg, "config.json")

    """

    def __init__(self, data: dict[str, Any] | None = None) -> None:
        """Initialize configuration with optional overrides.

        Parameters
        ----------
        data:
            Optional dictionary of configuration overrides.

        """
        merged = _deep_merge(DEFAULT_CONFIG, data or {})
        self._data: dict[str, Any] = merged
        # basic validation to catch obvious mistakes early
        validate_config(self)

    @property
    def optimization(self) -> _Section:
        """Return the optimization section of the configuration."""
        return _Section(self._data, ["optimization"])

    @property
    def model(self) -> _Section:
        """Return the model section of the configuration."""
        return _Section(self._data, ["model"])

    @property
    def conformer_generation(self) -> _Section:
        """Return the conformer generation section of the configuration."""
        return _Section(self._data, ["conformer_generation"])

    @property
    def technical(self) -> _Section:
        """Return the technical section of the configuration."""
        return _Section(self._data, ["technical"])

    def to_dict(self) -> dict[str, Any]:
        """Return a deep copy of the underlying configuration dictionary."""
        return copy.deepcopy(self._data)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Config:
        """Create a :class:`Config` instance from a plain dictionary."""
        return cls(data)

conformer_generation property

Return the conformer generation section of the configuration.

model property

Return the model section of the configuration.

optimization property

Return the optimization section of the configuration.

technical property

Return the technical section of the configuration.

__init__(data=None)

Initialize configuration with optional overrides.

Parameters

data: Optional dictionary of configuration overrides.

Source code in src/gpuma/config.py 
213
214
215
216
217
218
219
220
221
222
223
224
225
def __init__(self, data: dict[str, Any] | None = None) -> None:
    """Initialize configuration with optional overrides.

    Parameters
    ----------
    data:
        Optional dictionary of configuration overrides.

    """
    merged = _deep_merge(DEFAULT_CONFIG, data or {})
    self._data: dict[str, Any] = merged
    # basic validation to catch obvious mistakes early
    validate_config(self)

from_dict(data) classmethod

Create a :class:Config instance from a plain dictionary.

Source code in src/gpuma/config.py 
251
252
253
254
@classmethod
def from_dict(cls, data: dict[str, Any]) -> Config:
    """Create a :class:`Config` instance from a plain dictionary."""
    return cls(data)

to_dict()

Return a deep copy of the underlying configuration dictionary.

Source code in src/gpuma/config.py 
247
248
249
def to_dict(self) -> dict[str, Any]:
    """Return a deep copy of the underlying configuration dictionary."""
    return copy.deepcopy(self._data)

load_config_from_file(filepath='config.json')

Load configuration from a JSON/YAML file and deep-merge with defaults.

This function caches the raw dictionary loaded from the file to avoid repeated I/O and parsing. The returned Config object is always a new instance, safe to modify.

Parameters:

Name Type Description Default
filepath str

Path to the config file. If it doesn't exist, defaults are used.

 'config.json'

Returns:

Name Type Description
A Config

class:Config instance with data merged with :data:DEFAULT_CONFIG.
Config

Unknown keys are preserved.
Source code in src/gpuma/config.py 
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
def load_config_from_file(filepath: str = "config.json") -> Config:
    """Load configuration from a JSON/YAML file and deep-merge with defaults.

    This function caches the raw dictionary loaded from the file to avoid
    repeated I/O and parsing. The returned Config object is always a new
    instance, safe to modify.

    Args:
        filepath: Path to the config file. If it doesn't exist, defaults are used.

    Returns:
        A :class:`Config` instance with data merged with :data:`DEFAULT_CONFIG`.
        Unknown keys are preserved.

    """
    user_cfg = _read_config_file(filepath)

    if user_cfg is not None and not isinstance(user_cfg, dict):
        raise ValueError("Configuration file must contain a JSON/YAML object at the root")

    # Config.__init__ deep-merges via _deep_merge, so no need to deepcopy here
    cfg = Config.from_dict(user_cfg)
    validate_config(cfg)
    return cfg

save_config_to_file(config, filepath)

Save configuration to JSON/YAML file.

Accepts either a :class:Config instance or a plain dictionary.

Source code in src/gpuma/config.py 
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
def save_config_to_file(config: Any, filepath: str) -> None:
    """Save configuration to JSON/YAML file.

    Accepts either a :class:`Config` instance or a plain dictionary.
    """
    cfg_dict = config.to_dict() if isinstance(config, Config) else dict(config)

    with open(filepath, "w", encoding="utf-8") as f:
        if filepath.endswith(".json"):
            json.dump(cfg_dict, f, indent=2)
        elif filepath.endswith(".yaml") or filepath.endswith(".yml"):
            try:
                from yaml import safe_dump as _yaml_safe_dump  # type: ignore

                _yaml_safe_dump(cfg_dict, f, default_flow_style=False, sort_keys=False)
            except ImportError as e:
                raise ImportError("PyYAML is required to save YAML config files") from e
        else:
            raise ValueError("Config file must be JSON or YAML format")

    # Invalidate cache since the file on disk has changed
    _read_config_file.cache_clear()

resolve_model_type(config)

Normalize a model_type value to its canonical form.

Accepted aliases: - "fairchem" / "uma" -> "fairchem" - "orb" / "orb-v3" -> "orb" - "sevennet" / "7net" -> "sevennet"

Works with either :class:Config or a plain dict.

Source code in src/gpuma/config.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def resolve_model_type(config: Config | dict[str, Any]) -> str:
    """Normalize a ``model_type`` value to its canonical form.

    Accepted aliases:
    - ``"fairchem"`` / ``"uma"`` -> ``"fairchem"``
    - ``"orb"`` / ``"orb-v3"`` -> ``"orb"``
    - ``"sevennet"`` / ``"7net"`` -> ``"sevennet"``

    Works with either :class:`Config` or a plain dict.
    """
    if isinstance(config, Config):
        raw = str(getattr(config.model, "model_type", "fairchem")).strip().lower()
    else:
        raw = str(
            (config or {}).get("model", {}).get("model_type", "fairchem")
        ).strip().lower()
    canonical = _MODEL_TYPE_ALIASES.get(raw)
    if canonical is None:
        raise ValueError(
            f"Unknown model_type {raw!r}. Must be one of: {sorted(VALID_MODEL_TYPES)}"
        )
    return canonical

I/O & Structure Conversion

Functions for reading, writing, and converting molecular structures.

read_xyz(file_path, charge=0, multiplicity=1)

Read an XYZ file and return a :class:Structure instance.

Parameters

file_path: Path to the XYZ file to read. charge: Optional total charge to set on the structure (default: 0). multiplicity: Optional spin multiplicity to set (default: 1).

Returns

Structure Object with symbols, coordinates, and an optional comment.

Raises

FileNotFoundError If the specified file does not exist. ValueError If the file format is invalid.

Source code in src/gpuma/utils/io_handler.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def read_xyz(file_path: str, charge: int = 0, multiplicity: int = 1) -> Structure:
    """Read an XYZ file and return a :class:`Structure` instance.

    Parameters
    ----------
    file_path:
        Path to the XYZ file to read.
    charge:
        Optional total charge to set on the structure (default: ``0``).
    multiplicity:
        Optional spin multiplicity to set (default: ``1``).

    Returns
    -------
    Structure
        Object with symbols, coordinates, and an optional comment.

    Raises
    ------
    FileNotFoundError
        If the specified file does not exist.
    ValueError
        If the file format is invalid.

    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File {file_path} not found")

    symbols: list[str] = []
    coordinates: list[tuple[float, float, float]] = []

    try:
        with open(file_path, encoding="utf-8") as infile:
            # Read first line: number of atoms
            line = infile.readline()
            try:
                num_atoms = int(line.strip())
            except ValueError as exc:
                raise ValueError(
                    "First line must contain the number of atoms as an integer"
                ) from exc

            # Read second line: comment
            comment_line = infile.readline()
            # If EOF is reached, comment_line is "" (which is falsy)
            # But a blank line "\n" is truthy.
            if not comment_line and num_atoms >= 0:
                # We expected a comment line
                # Note: Original code calculated found as max(0, len(lines) - 2)
                # If we have 1 line, found = 0.
                raise ValueError(f"Expected {num_atoms} atom lines, but found 0")

            comment = comment_line.rstrip("\n")

            for i in range(num_atoms):
                line = infile.readline()
                if not line:
                    raise ValueError(f"Expected {num_atoms} atom lines, but found {i}")

                parts = line.split()
                if len(parts) < 4:
                    raise ValueError(f"Line {i + 3} must contain at least 4 elements: symbol x y z")
                symbol = parts[0]
                try:
                    x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
                except ValueError as exc:
                    raise ValueError(f"Invalid coordinates in line {i + 3}: {parts[1:4]}") from exc
                symbols.append(symbol)
                coordinates.append((x, y, z))

    except Exception as exc:
        if isinstance(exc, (FileNotFoundError, ValueError)):
            raise
        raise ValueError(f"Error reading XYZ file: {exc}") from exc

    return Structure(
        symbols=symbols,
        coordinates=coordinates,
        comment=comment,
        charge=charge,
        multiplicity=multiplicity,
    )

read_multi_xyz(file_path, charge=0, multiplicity=1)

Read an XYZ file containing multiple structures.

Parameters

file_path: Path to the multi-structure XYZ file. charge: Optional total charge to set on all returned structures (default: 0). multiplicity: Optional spin multiplicity to set (default: 1).

Returns

list[Structure] List of structures read from the file.

Raises

FileNotFoundError If the specified file does not exist. ValueError If the file format is invalid.

Source code in src/gpuma/utils/io_handler.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def read_multi_xyz(file_path: str, charge: int = 0, multiplicity: int = 1) -> list[Structure]:
    """Read an XYZ file containing multiple structures.

    Parameters
    ----------
    file_path:
        Path to the multi-structure XYZ file.
    charge:
        Optional total charge to set on all returned structures (default: ``0``).
    multiplicity:
        Optional spin multiplicity to set (default: ``1``).

    Returns
    -------
    list[Structure]
        List of structures read from the file.

    Raises
    ------
    FileNotFoundError
        If the specified file does not exist.
    ValueError
        If the file format is invalid.

    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File {file_path} not found")

    structures: list[Structure] = []

    try:
        with open(file_path, encoding="utf-8") as infile:
            line_iterator = iter(infile)
            while True:
                try:
                    line = next(line_iterator)
                except StopIteration:
                    break

                line_stripped = line.strip()
                if not line_stripped:
                    continue

                try:
                    num_atoms = int(line_stripped)
                except ValueError:
                    continue

                try:
                    comment_line = next(line_iterator)
                    comment = comment_line.rstrip("\n")
                except StopIteration:
                    break

                symbols: list[str] = []
                coordinates: list[tuple[float, float, float]] = []

                valid = True
                for atom_idx in range(num_atoms):
                    try:
                        atom_line = next(line_iterator)
                    except StopIteration:
                        logger.warning(
                            "Structure '%s': unexpected end of file at atom %d/%d, skipping",
                            comment, atom_idx + 1, num_atoms,
                        )
                        valid = False
                        break

                    parts = atom_line.split()
                    if len(parts) < 4:
                        logger.warning(
                            "Structure '%s': malformed atom line %d, skipping structure",
                            comment, atom_idx + 1,
                        )
                        valid = False
                        break

                    symbol = parts[0]
                    try:
                        x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
                    except ValueError:
                        logger.warning(
                            "Structure '%s': invalid coordinates at atom %d, skipping structure",
                            comment, atom_idx + 1,
                        )
                        valid = False
                        break
                    symbols.append(symbol)
                    coordinates.append((x, y, z))

                if valid and len(symbols) == num_atoms:
                    structures.append(
                        Structure(
                            symbols=symbols,
                            coordinates=coordinates,
                            comment=comment,
                            charge=charge,
                            multiplicity=multiplicity,
                        )
                    )

    except Exception as exc:
        raise ValueError(f"Error reading multi-XYZ file: {exc}") from exc

    return structures

read_xyz_directory(directory_path, charge=0, multiplicity=1, sort=True)

Read all XYZ files from a directory.

Parameters

directory_path: Path to directory containing XYZ files. charge: Optional total charge to set on all returned structures (default: 0). multiplicity: Optional spin multiplicity to set (default: 1). sort: Controls the order in which files are read. If True (default), files are sorted in natural numerical order (e.g. mol_2.xyz before mol_10.xyz). If False, files are read in OS-dependent order, which is not guaranteed and may vary between runs or platforms. If a callable is provided, it is used as the sort key instead (see :func:sorted).

Returns

list[Structure] List of structures from all XYZ files in the directory, in the order determined by sort.

Raises

FileNotFoundError If the directory does not exist. ValueError If no valid XYZ files are found.

Source code in src/gpuma/utils/io_handler.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
def read_xyz_directory(
    directory_path: str,
    charge: int = 0,
    multiplicity: int = 1,
    sort: bool | Callable[[str], list] = True,
) -> list[Structure]:
    """Read all XYZ files from a directory.

    Parameters
    ----------
    directory_path:
        Path to directory containing XYZ files.
    charge:
        Optional total charge to set on all returned structures (default: ``0``).
    multiplicity:
        Optional spin multiplicity to set (default: ``1``).
    sort:
        Controls the order in which files are read. If ``True`` (default),
        files are sorted in natural numerical order (e.g. ``mol_2.xyz``
        before ``mol_10.xyz``). If ``False``, files are read in
        OS-dependent order, which is not guaranteed and may vary between
        runs or platforms. If a callable is provided, it is used as the
        sort key instead (see :func:`sorted`).

    Returns
    -------
    list[Structure]
        List of structures from all XYZ files in the directory, in the
        order determined by ``sort``.

    Raises
    ------
    FileNotFoundError
        If the directory does not exist.
    ValueError
        If no valid XYZ files are found.

    """
    if not os.path.exists(directory_path):
        raise FileNotFoundError(f"Directory {directory_path} not found")

    xyz_files = glob.iglob(os.path.join(directory_path, "*.xyz"))

    if sort is True:
        xyz_files = sorted(xyz_files, key=natural_sort_key)
    elif callable(sort):
        xyz_files = sorted(xyz_files, key=sort)
    else:
        xyz_files = list(xyz_files)

    structures: list[Structure] = []
    found_any = False

    for xyz_file in xyz_files:
        found_any = True
        try:
            structures.append(read_xyz(xyz_file, charge=charge, multiplicity=multiplicity))
        except Exception as exc:  # pragma: no cover - logged and skipped
            logger.warning("Failed to read %s: %s", xyz_file, exc)

    if not found_any:
        raise ValueError(f"No XYZ files found in directory {directory_path}")

    if not structures:
        raise ValueError("No valid structures could be read from any XYZ files")

    return structures

smiles_to_xyz(smiles_string, return_full_xyz_str=False, multiplicity=None, config=None)

Convert a SMILES string to a :class:Structure or an XYZ string.

Parameters

smiles_string: Valid SMILES string representing the molecular structure. return_full_xyz_str: If True, return an XYZ-format string instead of a :class:Structure instance. multiplicity: Optional spin multiplicity to set on the structure (default: None).

Returns

Structure | str Either a :class:Structure or an XYZ string depending on return_full_xyz_str.

Source code in src/gpuma/utils/io_handler.py 
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
def smiles_to_xyz(
    smiles_string: str,
    return_full_xyz_str: bool = False,
    multiplicity: int | None = None,
    config: Config | None = None,
) -> Structure | str:
    """Convert a SMILES string to a :class:`Structure` or an XYZ string.

    Parameters
    ----------
    smiles_string:
        Valid SMILES string representing the molecular structure.
    return_full_xyz_str:
        If ``True``, return an XYZ-format string instead of a
        :class:`Structure` instance.
    multiplicity:
        Optional spin multiplicity to set on the structure (default: ``None``).

    Returns
    -------
    Structure | str
        Either a :class:`Structure` or an XYZ string depending on
        ``return_full_xyz_str``.

    """
    if not smiles_string or not smiles_string.strip():
        raise ValueError("SMILES string cannot be empty or None")

    struct = _smiles_to_structure_util(smiles_string.strip(), config=config)
    if multiplicity is not None:
        struct.multiplicity = int(multiplicity)

    if return_full_xyz_str:
        xyz_lines = [str(struct.n_atoms)]
        xyz_lines.append(
            f"Generated from SMILES using MORFEUS | "
            f"Charge: {struct.charge} | "
            f"Multiplicity: {struct.multiplicity}"
        )
        for atom, coord in zip(struct.symbols, struct.coordinates, strict=True):
            xyz_lines.append(f"{atom} {coord[0]:.6f} {coord[1]:.6f} {coord[2]:.6f}")
        return "\n".join(xyz_lines)

    struct.comment = (
        f"Generated from SMILES: {smiles_string} | "
        f"Charge: {struct.charge} | "
        f"Multiplicity: {struct.multiplicity}"
    )
    return struct

smiles_to_ensemble(smiles_string, max_num_confs, multiplicity=None, seed=None, config=None)

Generate conformer ensemble from SMILES.

Parameters

smiles_string: Valid SMILES string representing the molecular structure. max_num_confs: Maximum number of conformers to generate. multiplicity: Optional spin multiplicity to set on the structures (default: None). seed: Optional random seed for reproducible conformer generation.

Returns

list[Structure] A list of :class:Structure instances representing the conformers.

Source code in src/gpuma/utils/io_handler.py 
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
def smiles_to_ensemble(
    smiles_string: str,
    max_num_confs: int,
    multiplicity: int | None = None,
    seed: int | None = None,
    config: Config | None = None,
) -> list[Structure]:
    """Generate conformer ensemble from SMILES.

    Parameters
    ----------
    smiles_string:
        Valid SMILES string representing the molecular structure.
    max_num_confs:
        Maximum number of conformers to generate.
    multiplicity:
        Optional spin multiplicity to set on the structures (default: ``None``).
    seed:
        Optional random seed for reproducible conformer generation.

    Returns
    -------
    list[Structure]
        A list of :class:`Structure` instances representing the conformers.

    """
    if not smiles_string or not smiles_string.strip():
        raise ValueError("SMILES string cannot be empty or None")

    mult = int(multiplicity) if multiplicity is not None else 1
    structs = _smiles_to_ensemble_util(
        smiles_string.strip(), max_num_confs, multiplicity=mult, seed=seed, config=config,
    )
    return structs

save_xyz_file(structure, file_path)

Save a single :class:Structure to an XYZ file.

The comment line includes the energy (if set), charge, and multiplicity.

Parameters

structure: Structure to write. file_path: Destination file path.

Source code in src/gpuma/utils/io_handler.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
def save_xyz_file(structure: Structure, file_path: str) -> None:
    """Save a single :class:`Structure` to an XYZ file.

    The comment line includes the energy (if set), charge, and multiplicity.

    Parameters
    ----------
    structure:
        Structure to write.
    file_path:
        Destination file path.

    """
    lines: list[str] = [str(structure.n_atoms)]
    # include existing comment and ensure energy/charge/multiplicity are visible
    base_comment = structure.comment or ""
    energy_part = ""
    if structure.energy is not None:
        energy_part = f" | Energy: {structure.energy:.6f} eV"
    state_part = f" | Charge: {structure.charge} | Multiplicity: {structure.multiplicity}"
    comment = (base_comment + energy_part + state_part).strip() or "Structure"
    lines.append(comment)
    for symbol, coord in zip(structure.symbols, structure.coordinates, strict=True):
        lines.append(f"{symbol} {coord[0]:.6f} {coord[1]:.6f} {coord[2]:.6f}")

    parent = os.path.dirname(file_path)
    if parent:
        os.makedirs(parent, exist_ok=True)

    with open(file_path, "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines))
        fh.write("\n")

save_multi_xyz(structures, file_path, comments=None)

Save multiple structures to a single multi-structure XYZ file.

Each structure block includes the energy (if set), charge, and multiplicity in the comment line.

Parameters

structures: List of structures to write. file_path: Destination file path. comments: Optional per-structure comment strings. Falls back to each structure's own comment if not provided.

Source code in src/gpuma/utils/io_handler.py 
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
def save_multi_xyz(
    structures: list[Structure], file_path: str, comments: list[str] | None = None
) -> None:
    """Save multiple structures to a single multi-structure XYZ file.

    Each structure block includes the energy (if set), charge, and
    multiplicity in the comment line.

    Parameters
    ----------
    structures:
        List of structures to write.
    file_path:
        Destination file path.
    comments:
        Optional per-structure comment strings. Falls back to each
        structure's own comment if not provided.

    """
    lines: list[str] = []
    for idx, struct in enumerate(structures):
        lines.append(str(struct.n_atoms))
        base_comment = ""
        if comments and idx < len(comments):
            base_comment = comments[idx]
        elif struct.comment:
            base_comment = struct.comment
        energy_part = ""
        if struct.energy is not None:
            energy_part = f" | Energy: {struct.energy:.6f} eV"
        state_part = f" | Charge: {struct.charge} | Multiplicity: {struct.multiplicity}"
        comment = (base_comment + energy_part + state_part).strip() or f"Structure {idx + 1}"
        lines.append(comment)
        for symbol, coord in zip(struct.symbols, struct.coordinates, strict=True):
            lines.append(f"{symbol} {coord[0]:.6f} {coord[1]:.6f} {coord[2]:.6f}")

    parent = os.path.dirname(file_path)
    if parent:
        os.makedirs(parent, exist_ok=True)

    with open(file_path, "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines))
        fh.write("\n")

save_as_single_xyz_files(structures, output_dir, comments=None)

Save each structure to its own XYZ file in a directory.

Files are zero-padded to sort naturally, e.g. structure_01.xyz for up to 99 structures, structure_0001.xyz for up to 9999, etc.

Parameters

structures: List of structures to save. output_dir: Directory where files will be written. Created if it does not exist. comments: Optional per-structure comment strings.

Source code in src/gpuma/utils/io_handler.py 
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
def save_as_single_xyz_files(
    structures: list[Structure], output_dir: str, comments: list[str] | None = None
) -> None:
    """Save each structure to its own XYZ file in a directory.

    Files are zero-padded to sort naturally, e.g. ``structure_01.xyz`` for
    up to 99 structures, ``structure_0001.xyz`` for up to 9999, etc.

    Parameters
    ----------
    structures:
        List of structures to save.
    output_dir:
        Directory where files will be written. Created if it does not exist.
    comments:
        Optional per-structure comment strings.

    """
    os.makedirs(output_dir, exist_ok=True)
    width = len(str(len(structures)))
    for idx, struct in enumerate(structures):
        if comments and idx < len(comments):
            struct = Structure(
                symbols=struct.symbols,
                coordinates=struct.coordinates,
                energy=struct.energy,
                charge=struct.charge,
                multiplicity=struct.multiplicity,
                comment=comments[idx],
            )
        file_path = os.path.join(output_dir, f"structure_{idx + 1:0{width}d}.xyz")
        save_xyz_file(struct, file_path)

Batched Conformer Generation

Convert many SMILES to 3D structures in one call. These batch across molecules, which is what makes the optional GPU backend worthwhile; the per-molecule helpers above call into them. The backend follows technical.device in the configuration and falls back to CPU when no GPU is usable.

generate_structures(smiles_list, config=None, multiplicity=None, n_confs=None, seed=DEFAULT_SEED, prune_rms_thresh=DEFAULT_PRUNE_RMS, max_iters=DEFAULT_MAX_ITERS, batch_size=500, n_threads=1, allow_cpu_fallback=True)

Convert SMILES to one 3D structure each, as a single batch.

Embeds a conformer ensemble per molecule, minimizes every conformer, and keeps the lowest-energy one. Use :func:generate_ensembles to keep more than one.

Parameters

smiles_list: SMILES strings to convert. config: gpuma configuration; technical.device selects the backend. Loaded from the default location if omitted. multiplicity: Spin multiplicity applied to every returned structure. None takes config.optimization.multiplicity. n_confs: Conformers generated per molecule. None selects a budget from rotatable-bond count via :data:CONF_BUDGET; an explicit int applies that count uniformly instead. seed: Random seed for the embedding. -1 lets RDKit choose one per run, making geometries non-reproducible. prune_rms_thresh: RMSD threshold for discarding duplicate conformers during embedding. max_iters: Maximum force-field minimization iterations per conformer. GPU backend only -- the CPU backend goes through morfeus, which does not expose it. batch_size: Molecules per nvMolKit batch. GPU backend only. n_threads: Thread count handed to RDKit, CPU backend only. Leave at 1 when the caller already parallelizes across molecules. allow_cpu_fallback: When a GPU was requested but is unusable, True transparently runs on CPU. Set False to re-raise instead, so a caller that parallelizes CPU work itself can choose its own strategy rather than silently getting a serial CPU run.

Returns

list[Structure | None] One entry per input, in input order. None marks a molecule that could not be parsed or embedded, so one bad SMILES cannot abort a library.

Raises

Exception Whatever the GPU backend raised, when a GPU was requested and allow_cpu_fallback is False.

Source code in src/gpuma/conformer_generation/embed.py 
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
def generate_structures(
    smiles_list: list[str],
    config: Config | None = None,
    multiplicity: int | None = None,
    n_confs: int | None = None,
    seed: int = DEFAULT_SEED,
    prune_rms_thresh: float = DEFAULT_PRUNE_RMS,
    max_iters: int = DEFAULT_MAX_ITERS,
    batch_size: int = 500,
    n_threads: int = 1,
    allow_cpu_fallback: bool = True,
) -> list[Structure | None]:
    """Convert SMILES to one 3D structure each, as a single batch.

    Embeds a conformer ensemble per molecule, minimizes every conformer, and
    keeps the lowest-energy one. Use :func:`generate_ensembles` to keep more
    than one.

    Parameters
    ----------
    smiles_list:
        SMILES strings to convert.
    config:
        gpuma configuration; ``technical.device`` selects the backend. Loaded
        from the default location if omitted.
    multiplicity:
        Spin multiplicity applied to every returned structure. ``None`` takes
        ``config.optimization.multiplicity``.
    n_confs:
        Conformers generated per molecule. ``None`` selects a budget from
        rotatable-bond count via :data:`CONF_BUDGET`; an explicit int applies
        that count uniformly instead.
    seed:
        Random seed for the embedding. ``-1`` lets RDKit choose one per run,
        making geometries non-reproducible.
    prune_rms_thresh:
        RMSD threshold for discarding duplicate conformers during embedding.
    max_iters:
        Maximum force-field minimization iterations per conformer. GPU backend
        only -- the CPU backend goes through morfeus, which does not expose it.
    batch_size:
        Molecules per nvMolKit batch. GPU backend only.
    n_threads:
        Thread count handed to RDKit, CPU backend only. Leave at 1 when the
        caller already parallelizes across molecules.
    allow_cpu_fallback:
        When a GPU was requested but is unusable, ``True`` transparently runs
        on CPU. Set ``False`` to re-raise instead, so a caller that
        parallelizes CPU work itself can choose its own strategy rather than
        silently getting a serial CPU run.

    Returns
    -------
    list[Structure | None]
        One entry per input, in input order. ``None`` marks a molecule that
        could not be parsed or embedded, so one bad SMILES cannot abort a
        library.

    Raises
    ------
    Exception
        Whatever the GPU backend raised, when a GPU was requested and
        ``allow_cpu_fallback`` is ``False``.
    """
    batches = _generate(
        smiles_list,
        config,
        multiplicity,
        n_confs,
        seed,
        prune_rms_thresh,
        max_iters,
        batch_size,
        n_threads,
        allow_cpu_fallback,
        n_keep=1,
    )
    return [structures[0] if structures else None for structures in batches]

generate_ensembles(smiles_list, max_num_confs, config=None, multiplicity=None, n_confs=None, seed=DEFAULT_SEED, prune_rms_thresh=DEFAULT_PRUNE_RMS, max_iters=DEFAULT_MAX_ITERS, batch_size=500, n_threads=1, allow_cpu_fallback=True)

Convert SMILES to conformer ensembles, as a single batch.

As :func:generate_structures, but keeps several conformers per molecule instead of one. All other parameters carry the same meaning.

Parameters

max_num_confs: Maximum conformers returned per molecule, lowest energy first. This is distinct from n_confs, which controls how many are generated -- generating fewer than you keep simply wastes the budget. Fewer than requested may come back either way, since RMSD pruning removes duplicates.

Returns

list[list[Structure] | None] One entry per input, in input order. Each entry is a list of at most max_num_confs structures, or None for a molecule that could not be parsed or embedded.

Raises

ValueError If max_num_confs is not positive. Exception Whatever the GPU backend raised, when a GPU was requested and allow_cpu_fallback is False.

Source code in src/gpuma/conformer_generation/embed.py 
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def generate_ensembles(
    smiles_list: list[str],
    max_num_confs: int,
    config: Config | None = None,
    multiplicity: int | None = None,
    n_confs: int | None = None,
    seed: int = DEFAULT_SEED,
    prune_rms_thresh: float = DEFAULT_PRUNE_RMS,
    max_iters: int = DEFAULT_MAX_ITERS,
    batch_size: int = 500,
    n_threads: int = 1,
    allow_cpu_fallback: bool = True,
) -> list[list[Structure] | None]:
    """Convert SMILES to conformer ensembles, as a single batch.

    As :func:`generate_structures`, but keeps several conformers per molecule
    instead of one. All other parameters carry the same meaning.

    Parameters
    ----------
    max_num_confs:
        Maximum conformers *returned* per molecule, lowest energy first. This
        is distinct from ``n_confs``, which controls how many are *generated* --
        generating fewer than you keep simply wastes the budget. Fewer than
        requested may come back either way, since RMSD pruning removes
        duplicates.

    Returns
    -------
    list[list[Structure] | None]
        One entry per input, in input order. Each entry is a list of at most
        ``max_num_confs`` structures, or ``None`` for a molecule that could not
        be parsed or embedded.

    Raises
    ------
    ValueError
        If ``max_num_confs`` is not positive.
    Exception
        Whatever the GPU backend raised, when a GPU was requested and
        ``allow_cpu_fallback`` is ``False``.
    """
    if max_num_confs <= 0:
        raise ValueError(f"max_num_confs must be positive, got {max_num_confs}")
    return _generate(
        smiles_list,
        config,
        multiplicity,
        n_confs,
        seed,
        prune_rms_thresh,
        max_iters,
        batch_size,
        n_threads,
        allow_cpu_fallback,
        n_keep=max_num_confs,
    )

Low-Level Optimization

Lower-level functions used by the high-level API.

optimize_single_structure(structure, config=None, calculator=None)

Optimize a single :class:Structure using an ASE optimizer.

The same structure instance is returned with updated coordinates and energy.

Parameters

structure : Structure Molecular structure to optimize. config : Config, optional Configuration controlling the model and convergence settings. Defaults to :func:load_config_from_file if not provided. calculator : optional Pre-loaded ASE calculator. If None, one is loaded (and cached) from the configuration.

Returns

Structure The input structure with optimized coordinates and energy set.

Raises

RuntimeError If the optimization fails for any reason.

Source code in src/gpuma/optimizer.py 
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def optimize_single_structure(
    structure: Structure,
    config: Config | None = None,
    calculator: Any | None = None,
) -> Structure:
    """Optimize a single :class:`Structure` using an ASE optimizer.

    The same ``structure`` instance is returned with updated coordinates and
    energy.

    Parameters
    ----------
    structure : Structure
        Molecular structure to optimize.
    config : Config, optional
        Configuration controlling the model and convergence settings.
        Defaults to :func:`load_config_from_file` if not provided.
    calculator : optional
        Pre-loaded ASE calculator. If ``None``, one is loaded (and cached)
        from the configuration.

    Returns
    -------
    Structure
        The input structure with optimized coordinates and energy set.

    Raises
    ------
    RuntimeError
        If the optimization fails for any reason.
    """
    if config is None:
        config = load_config_from_file()

    try:
        if calculator is None:
            calculator = _get_cached_calculator(config)

        atoms = Atoms(symbols=structure.symbols, positions=structure.coordinates)
        atoms.calc = calculator
        atoms.info = {"charge": structure.charge, "spin": structure.multiplicity}

        fmax = _resolve_force_criterion(config)
        opt_cls, opt_name = _resolve_ase_optimizer(config)

        logger.info(
            "Starting single geometry optimization for structure with %d atoms "
            "(optimizer=%s)",
            structure.n_atoms,
            opt_name,
        )
        optimizer = opt_cls(atoms, logfile=None)
        optimizer.run(fmax=fmax)
        logger.info("Optimization completed after %d steps", optimizer.get_number_of_steps())

        structure.coordinates = atoms.get_positions().tolist()
        structure.energy = float(atoms.get_potential_energy())
        return structure

    except Exception as exc:  # pragma: no cover - defensive logging
        raise RuntimeError(f"Optimization failed: {exc}") from exc

optimize_structure_batch(structures, config=None)

Optimize a list of structures and return them with updated coordinates.

The optimization mode is controlled by config.optimization.batch_optimization_mode:

  • "sequential": Each structure is optimized individually with ASE using a shared calculator.
  • "batch": All structures are optimized together using torch-sim GPU-accelerated batch optimization (requires GPU).
Parameters

structures : list[Structure] Structures to optimize. config : Config, optional Configuration object. Defaults to :func:load_config_from_file.

Returns

list[Structure] Optimized structures with coordinates and energies set.

Raises

ValueError If structures have mismatched symbols/coordinates or are empty, or if the optimization mode is unknown.

Source code in src/gpuma/optimizer.py 
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
def optimize_structure_batch(
    structures: list[Structure],
    config: Config | None = None,
) -> list[Structure]:
    """Optimize a list of structures and return them with updated coordinates.

    The optimization mode is controlled by
    ``config.optimization.batch_optimization_mode``:

    - ``"sequential"``: Each structure is optimized individually with
      ASE using a shared calculator.
    - ``"batch"``: All structures are optimized together using torch-sim
      GPU-accelerated batch optimization (requires GPU).

    Parameters
    ----------
    structures : list[Structure]
        Structures to optimize.
    config : Config, optional
        Configuration object. Defaults to :func:`load_config_from_file`.

    Returns
    -------
    list[Structure]
        Optimized structures with coordinates and energies set.

    Raises
    ------
    ValueError
        If structures have mismatched symbols/coordinates or are empty,
        or if the optimization mode is unknown.
    """
    if config is None:
        config = load_config_from_file()
    if not structures:
        return []

    for i, struct in enumerate(structures):
        if struct.n_atoms != len(struct.coordinates):
            raise ValueError(f"Structure {i}: symbols/coords length mismatch")
        if struct.n_atoms == 0:
            raise ValueError(f"Structure {i}: empty structure")

    on_cpu = _parse_device_string(config.technical.device) == "cpu"
    mode = str(config.optimization.batch_optimization_mode).lower()

    logger.info("Optimization device: %s", "CPU" if on_cpu else "GPU")

    with timed_block("Total optimization") as tb:
        if mode == "sequential" or on_cpu:
            if not on_cpu and mode == "batch":
                logger.warning(
                    "Batch optimization mode requires GPU, falling back to sequential mode on CPU.",
                )
            results = _optimize_sequential(structures, config)
        elif mode == "batch":
            results = _optimize_batch(structures, config)
        else:
            raise ValueError(
                f"Unknown optimization mode: {mode!r}. Use 'sequential' or 'batch' "
                "(batch requires GPU)."
            )

    log_optimization_summary(structures, results, tb.elapsed, mode, config)
    return results

Utilities

Timing decorators and context managers for profiling.

time_it(func)

Measure the execution time of a function and log the result.

Parameters

func: Callable to be wrapped.

Returns

callable Wrapped function that logs its runtime at :mod:logging.INFO level.

Source code in src/gpuma/utils/decorators.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def time_it(func):
    """Measure the execution time of a function and log the result.

    Parameters
    ----------
    func:
        Callable to be wrapped.

    Returns
    -------
    callable
        Wrapped function that logs its runtime at :mod:`logging.INFO` level.

    """

    @wraps(func)
    def wrap(*args, **kwargs):
        start_time = perf_counter()
        result = func(*args, **kwargs)
        elapsed = perf_counter() - start_time
        logger.info("Function: %r took: %.2f sec", func.__name__, elapsed)
        return result

    return wrap

timed_block

Context manager that measures and logs a named code block.

The elapsed time (in seconds) is available via the :attr:elapsed attribute after the block exits. While inside any :class:capture_timings context manager, the (name, elapsed) pair is also published to that capture.

Example

with timed_block("model loading") as tb: ... model = load_model() print(tb.elapsed)

Source code in src/gpuma/utils/decorators.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class timed_block:
    """Context manager that measures and logs a named code block.

    The elapsed time (in seconds) is available via the :attr:`elapsed`
    attribute after the block exits. While inside any
    :class:`capture_timings` context manager, the (name, elapsed) pair
    is also published to that capture.

    Example
    -------
    >>> with timed_block("model loading") as tb:
    ...     model = load_model()
    >>> print(tb.elapsed)
    """

    def __init__(self, name: str, *, level: int = logging.INFO):
        self.name = name
        self.elapsed: float = 0.0
        self._level = level

    def __enter__(self):
        """Start the timer and return ``self`` for attribute access."""
        self._start = perf_counter()
        return self

    def __exit__(self, *exc_info):
        """Stop the timer, store elapsed time, log, and notify captures."""
        self.elapsed = perf_counter() - self._start
        logger.log(self._level, "%s took %.2f sec", self.name, self.elapsed)
        for capture in _active_captures.get():
            capture._record(self.name, self.elapsed)

__enter__()

Start the timer and return self for attribute access.

Source code in src/gpuma/utils/decorators.py 
71
72
73
74
def __enter__(self):
    """Start the timer and return ``self`` for attribute access."""
    self._start = perf_counter()
    return self

__exit__(*exc_info)

Stop the timer, store elapsed time, log, and notify captures.

Source code in src/gpuma/utils/decorators.py 
76
77
78
79
80
81
def __exit__(self, *exc_info):
    """Stop the timer, store elapsed time, log, and notify captures."""
    self.elapsed = perf_counter() - self._start
    logger.log(self._level, "%s took %.2f sec", self.name, self.elapsed)
    for capture in _active_captures.get():
        capture._record(self.name, self.elapsed)