Skip to content

API Reference

This section contains the auto-generated API documentation for the toolbox package, created from the numpy-style docstrings in the source code.

Calculator Module

Bash Calculator

toolbox.calculator.bash

Bash calculator for running external commands.

This module provides a calculator interface for running bash commands and external programs with proper error handling and logging.

BashCalculator

A calculator for running bash commands in a specified working directory.

This class provides functionality to execute shell commands with proper directory management and logging.

Source code in toolbox/calculator/bash.py 
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class BashCalculator:
    """A calculator for running bash commands in a specified working directory.

    This class provides functionality to execute shell commands with proper
    directory management and logging.
    """

    def __init__(self, work_dir) -> None:
        """Initialize the BashCalculator.

        Parameters
        ----------
        work_dir : str
            The working directory where commands will be executed
        """
        self.work_dir = work_dir

    def run(
        self,
        command: str = None,
        stdin: str = None,
        stdout: str = "job.stdout",
        stderr: str = "job.stderr",
        mpi_command: str = "mpiexec.hydra",
        ignore_finished_tag: bool = False,
    ):
        """Execute a command in the working directory.

        Parameters
        ----------
        command : str, optional
            Command to execute, by default None
        stdin : str, optional
            Standard input file, by default None
        stdout : str, optional
            Standard output file, by default "job.stdout"
        stderr : str, optional
            Standard error file, by default "job.stderr"
        mpi_command : str, optional
            MPI command to use, by default "mpiexec.hydra"
        ignore_finished_tag : bool, optional
            Whether to ignore existing finished tag, by default False
        """
        if (ignore_finished_tag) or (
            not os.path.exists(os.path.join(self.work_dir, "finished_tag"))
        ):
            if mpi_command is not None:
                command = f"{mpi_command} {command}"

            root_dir = os.getcwd()
            os.chdir(self.work_dir)
            logging.info("{:=^50}".format(" Start calculation "))

            logging.info(f"Path: {os.getcwd()}")
            if stdin is not None:
                command += f" {stdin} "
            if stdout is not None:
                command += f" 1> {stdout} "
            if stderr is not None:
                command += f" 2> {stderr} "
            os.system(command=command)

            self._make_finished_tag(stdout)

            logging.info("{:=^50}".format(" End calculation "))
            os.chdir(root_dir)

    @staticmethod
    def _make_finished_tag(stdout):
        """Create a finished tag file after successful execution.

        Parameters
        ----------
        stdout : str
            Name of the standard output file to check
        """
        pass
__init__(work_dir)

Initialize the BashCalculator.

Parameters:

Name Type Description Default
work_dir str

The working directory where commands will be executed

 required
Source code in toolbox/calculator/bash.py 
19
20
21
22
23
24
25
26
27
def __init__(self, work_dir) -> None:
    """Initialize the BashCalculator.

    Parameters
    ----------
    work_dir : str
        The working directory where commands will be executed
    """
    self.work_dir = work_dir
run(command=None, stdin=None, stdout='job.stdout', stderr='job.stderr', mpi_command='mpiexec.hydra', ignore_finished_tag=False)

Execute a command in the working directory.

Parameters:

Name Type Description Default
command str

Command to execute, by default None

 None
stdin str

Standard input file, by default None

 None
stdout str

Standard output file, by default "job.stdout"

 'job.stdout'
stderr str

Standard error file, by default "job.stderr"

 'job.stderr'
mpi_command str

MPI command to use, by default "mpiexec.hydra"

 'mpiexec.hydra'
ignore_finished_tag bool

Whether to ignore existing finished tag, by default False

 False
Source code in toolbox/calculator/bash.py 
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
70
71
72
73
74
75
76
77
def run(
    self,
    command: str = None,
    stdin: str = None,
    stdout: str = "job.stdout",
    stderr: str = "job.stderr",
    mpi_command: str = "mpiexec.hydra",
    ignore_finished_tag: bool = False,
):
    """Execute a command in the working directory.

    Parameters
    ----------
    command : str, optional
        Command to execute, by default None
    stdin : str, optional
        Standard input file, by default None
    stdout : str, optional
        Standard output file, by default "job.stdout"
    stderr : str, optional
        Standard error file, by default "job.stderr"
    mpi_command : str, optional
        MPI command to use, by default "mpiexec.hydra"
    ignore_finished_tag : bool, optional
        Whether to ignore existing finished tag, by default False
    """
    if (ignore_finished_tag) or (
        not os.path.exists(os.path.join(self.work_dir, "finished_tag"))
    ):
        if mpi_command is not None:
            command = f"{mpi_command} {command}"

        root_dir = os.getcwd()
        os.chdir(self.work_dir)
        logging.info("{:=^50}".format(" Start calculation "))

        logging.info(f"Path: {os.getcwd()}")
        if stdin is not None:
            command += f" {stdin} "
        if stdout is not None:
            command += f" 1> {stdout} "
        if stderr is not None:
            command += f" 2> {stderr} "
        os.system(command=command)

        self._make_finished_tag(stdout)

        logging.info("{:=^50}".format(" End calculation "))
        os.chdir(root_dir)

Coulomb Calculator

toolbox.calculator.coulomb

Coulomb interaction calculators for molecular simulations.

This module provides calculators for: - Coulomb cutoff interactions - Ewald summation methods - Particle Mesh Ewald (PME) calculations - D3 dispersion corrections

These calculators are used for computing electrostatic interactions in molecular dynamics simulations.

CoulCutCalculator

Calculator for Coulomb interactions with cutoff.

This calculator computes Coulomb interactions between atoms using a simple cutoff scheme.

Source code in toolbox/calculator/coulomb.py 
107
108
109
110
111
112
113
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
class CoulCutCalculator:
    """Calculator for Coulomb interactions with cutoff.

    This calculator computes Coulomb interactions between atoms
    using a simple cutoff scheme.
    """

    def __init__(self, cutoff=5.0) -> None:
        """Initialize CoulCutCalculator.

        Parameters
        ----------
        cutoff : float, optional
            Cutoff distance in Angstroms, by default 5.0
        """
        self.cutoff = cutoff

    def calculate(self, atoms):
        """Calculate Coulomb energy and forces.

        Parameters
        ----------
        atoms : ase.Atoms
            ASE Atoms object with positions and charges

        Returns
        -------
        tuple
            Tuple of (energy, forces) where energy is in eV
            and forces is an array in eV/A
        """
        cellpar = atoms.cell.cellpar()
        coords = atoms.get_positions()
        charges = atoms.get_initial_charges()
        dist_mat = distance_array(coords, coords, box=cellpar)
        nat = len(atoms)
        forces = np.zeros((nat, 3))
        energy = 0.0
        for ii in range(nat):
            force_mask = (dist_mat[ii] < self.cutoff) & (np.arange(nat) > ii)
            sel_ids = np.where(force_mask)[0]
            forcecoul = qqrd2e * charges[ii] * charges[sel_ids] / dist_mat[ii][sel_ids]
            fpair = forcecoul / dist_mat[ii][sel_ids] ** 2
            delta_x = minimize_vectors(
                coords[ii].reshape(1, 3) - coords[sel_ids], box=cellpar
            )
            forces[ii] += np.sum(delta_x * fpair.reshape(-1, 1), axis=0)
            forces[sel_ids] -= delta_x * fpair.reshape(-1, 1)
            energy += np.sum(forcecoul)
        return energy, forces
__init__(cutoff=5.0)

Initialize CoulCutCalculator.

Parameters:

Name Type Description Default
cutoff float

Cutoff distance in Angstroms, by default 5.0

 5.0
Source code in toolbox/calculator/coulomb.py 
114
115
116
117
118
119
120
121
122
def __init__(self, cutoff=5.0) -> None:
    """Initialize CoulCutCalculator.

    Parameters
    ----------
    cutoff : float, optional
        Cutoff distance in Angstroms, by default 5.0
    """
    self.cutoff = cutoff
calculate(atoms)

Calculate Coulomb energy and forces.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object with positions and charges

 required

Returns:

Type Description
tuple

Tuple of (energy, forces) where energy is in eV and forces is an array in eV/A
Source code in toolbox/calculator/coulomb.py 
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
def calculate(self, atoms):
    """Calculate Coulomb energy and forces.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object with positions and charges

    Returns
    -------
    tuple
        Tuple of (energy, forces) where energy is in eV
        and forces is an array in eV/A
    """
    cellpar = atoms.cell.cellpar()
    coords = atoms.get_positions()
    charges = atoms.get_initial_charges()
    dist_mat = distance_array(coords, coords, box=cellpar)
    nat = len(atoms)
    forces = np.zeros((nat, 3))
    energy = 0.0
    for ii in range(nat):
        force_mask = (dist_mat[ii] < self.cutoff) & (np.arange(nat) > ii)
        sel_ids = np.where(force_mask)[0]
        forcecoul = qqrd2e * charges[ii] * charges[sel_ids] / dist_mat[ii][sel_ids]
        fpair = forcecoul / dist_mat[ii][sel_ids] ** 2
        delta_x = minimize_vectors(
            coords[ii].reshape(1, 3) - coords[sel_ids], box=cellpar
        )
        forces[ii] += np.sum(delta_x * fpair.reshape(-1, 1), axis=0)
        forces[sel_ids] -= delta_x * fpair.reshape(-1, 1)
        energy += np.sum(forcecoul)
    return energy, forces

CoulLongCalculator

Calculator for long-range Coulomb interactions with Ewald summation.

This calculator computes Coulomb interactions using Ewald summation with complementary error function screening.

Source code in toolbox/calculator/coulomb.py 
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
class CoulLongCalculator:
    """Calculator for long-range Coulomb interactions with Ewald summation.

    This calculator computes Coulomb interactions using Ewald summation
    with complementary error function screening.
    """

    def __init__(self, gewald, cutoff=5.0) -> None:
        """Initialize CoulLongCalculator.

        Parameters
        ----------
        gewald : float
            Ewald screening parameter
        cutoff : float, optional
            Real-space cutoff distance in Angstroms, by default 5.0
        """
        self.cutoff = cutoff
        self.gewald = gewald

    def calculate(self, atoms):
        """Calculate Coulomb energy and forces with Ewald summation.

        Parameters
        ----------
        atoms : ase.Atoms
            ASE Atoms object with positions and charges

        Returns
        -------
        tuple
            Tuple of (energy, forces) where energy is in eV
            and forces is an array in eV/A
        """
        cellpar = atoms.cell.cellpar()
        coords = atoms.get_positions()
        charges = atoms.get_initial_charges()
        dist_mat = distance_array(coords, coords, box=cellpar)
        nat = len(atoms)
        forces = np.zeros((nat, 3))
        energy = 0.0
        for ii in range(nat):
            force_mask = (dist_mat[ii] < self.cutoff) & (np.arange(nat) > ii)
            sel_ids = np.where(force_mask)[0]
            prefactor = qqrd2e * charges[ii] * charges[sel_ids] / dist_mat[ii][sel_ids]
            grij = self.gewald * dist_mat[ii][sel_ids]
            expm2 = np.exp(-grij * grij)
            t = 1.0 / (1.0 + EWALD_P * grij)
            erfc = t * (A1 + t * (A2 + t * (A3 + t * (A4 + t * A5)))) * expm2
            forcecoul = prefactor * (erfc + EWALD_F * grij * expm2)
            fpair = forcecoul / dist_mat[ii][sel_ids] ** 2
            delta_x = minimize_vectors(
                coords[ii].reshape(1, 3) - coords[sel_ids], box=cellpar
            )
            forces[ii] += np.sum(delta_x * fpair.reshape(-1, 1), axis=0)
            forces[sel_ids] -= delta_x * fpair.reshape(-1, 1)
            energy += np.sum(prefactor * erfc)
        return energy, forces
__init__(gewald, cutoff=5.0)

Initialize CoulLongCalculator.

Parameters:

Name Type Description Default
gewald float

Ewald screening parameter

 required
cutoff float

Real-space cutoff distance in Angstroms, by default 5.0

 5.0
Source code in toolbox/calculator/coulomb.py 
166
167
168
169
170
171
172
173
174
175
176
177
def __init__(self, gewald, cutoff=5.0) -> None:
    """Initialize CoulLongCalculator.

    Parameters
    ----------
    gewald : float
        Ewald screening parameter
    cutoff : float, optional
        Real-space cutoff distance in Angstroms, by default 5.0
    """
    self.cutoff = cutoff
    self.gewald = gewald
calculate(atoms)

Calculate Coulomb energy and forces with Ewald summation.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object with positions and charges

 required

Returns:

Type Description
tuple

Tuple of (energy, forces) where energy is in eV and forces is an array in eV/A
Source code in toolbox/calculator/coulomb.py 
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
def calculate(self, atoms):
    """Calculate Coulomb energy and forces with Ewald summation.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object with positions and charges

    Returns
    -------
    tuple
        Tuple of (energy, forces) where energy is in eV
        and forces is an array in eV/A
    """
    cellpar = atoms.cell.cellpar()
    coords = atoms.get_positions()
    charges = atoms.get_initial_charges()
    dist_mat = distance_array(coords, coords, box=cellpar)
    nat = len(atoms)
    forces = np.zeros((nat, 3))
    energy = 0.0
    for ii in range(nat):
        force_mask = (dist_mat[ii] < self.cutoff) & (np.arange(nat) > ii)
        sel_ids = np.where(force_mask)[0]
        prefactor = qqrd2e * charges[ii] * charges[sel_ids] / dist_mat[ii][sel_ids]
        grij = self.gewald * dist_mat[ii][sel_ids]
        expm2 = np.exp(-grij * grij)
        t = 1.0 / (1.0 + EWALD_P * grij)
        erfc = t * (A1 + t * (A2 + t * (A3 + t * (A4 + t * A5)))) * expm2
        forcecoul = prefactor * (erfc + EWALD_F * grij * expm2)
        fpair = forcecoul / dist_mat[ii][sel_ids] ** 2
        delta_x = minimize_vectors(
            coords[ii].reshape(1, 3) - coords[sel_ids], box=cellpar
        )
        forces[ii] += np.sum(delta_x * fpair.reshape(-1, 1), axis=0)
        forces[sel_ids] -= delta_x * fpair.reshape(-1, 1)
        energy += np.sum(prefactor * erfc)
    return energy, forces

DMFFPMECalculator

Calculator for Coulomb interactions using Particle Mesh Ewald (PME).

This calculator uses the DMFF library to compute Coulomb interactions with PME for efficient long-range electrostatics.

Source code in toolbox/calculator/coulomb.py 
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
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
300
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
class DMFFPMECalculator:
    """Calculator for Coulomb interactions using Particle Mesh Ewald (PME).

    This calculator uses the DMFF library to compute Coulomb interactions
    with PME for efficient long-range electrostatics.
    """

    def __init__(self, cutoff=6.0) -> None:
        """Initialize DMFFPMECalculator.

        Parameters
        ----------
        cutoff : float, optional
            Real-space cutoff distance in Angstroms, by default 6.0
        """
        self.cutoff = cutoff

    def calculate(self, atoms, ethresh=1e-6) -> float:
        """Calculate Coulomb energy using PME.

        Parameters
        ----------
        atoms : ase.Atoms
            ASE Atoms object with positions and charges
        ethresh : float, optional
            Ewald threshold, by default 1e-6

        Returns
        -------
        float
            Coulomb energy in eV/particle
        """
        positions = atoms.get_positions()
        box = atoms.get_cell()
        charges = atoms.get_initial_charges()
        pairs = calculate_pairs(atoms, self.cutoff)

        box = torch.tensor(box, dtype=torch.double, requires_grad=False, device=DEVICE)
        positions = torch.tensor(
            positions, dtype=torch.double, requires_grad=True, device=DEVICE
        )
        charges = torch.tensor(
            charges, dtype=torch.double, requires_grad=False, device=DEVICE
        )
        pairs = torch.tensor(
            pairs, dtype=torch.int32, requires_grad=False, device=DEVICE
        )

        mscales = torch.tensor([0.0, 0.0, 1.0, 1.0, 1.0, 1.0], device=DEVICE)
        kappa, K = self.setup_ewald(box, ethresh)
        K1, K2, K3 = K
        pme_recip_fn = generate_pme_recip(
            Ck_fn=Ck_1,
            kappa=kappa,
            gamma=False,
            pme_order=6,
            K1=K1,
            K2=K2,
            K3=K3,
            lmax=0,
        )
        charges = torch.reshape(charges, (-1, 1))
        energy = energy_pme(
            positions,
            box,
            pairs,
            charges,
            None,
            None,
            None,
            mscales,
            None,
            None,
            None,
            pme_recip_fn,
            kappa,
            K1,
            K2,
            K3,
            0,
            False,
            True,
        )[0]
        # from kJ/mol to eV/particle
        return (energy * ENERGY_COEFF).item()

    def setup_ewald(self, box, ethresh):
        """Set up Ewald parameters for PME calculation.

        Parameters
        ----------
        box : array_like
            Simulation box vectors
        ethresh : float
            Ewald threshold

        Returns
        -------
        tuple
            Tuple of (kappa, K) where kappa is the screening parameter
            and K is the grid dimensions
        """
        kappa, K1, K2, K3 = setup_ewald_parameters(
            torch.tensor(self.cutoff), torch.tensor(ethresh), box, 0.01, "openmm"
        )
        K = (K1, K2, K3)
        return kappa, K
__init__(cutoff=6.0)

Initialize DMFFPMECalculator.

Parameters:

Name Type Description Default
cutoff float

Real-space cutoff distance in Angstroms, by default 6.0

 6.0
Source code in toolbox/calculator/coulomb.py 
226
227
228
229
230
231
232
233
234
def __init__(self, cutoff=6.0) -> None:
    """Initialize DMFFPMECalculator.

    Parameters
    ----------
    cutoff : float, optional
        Real-space cutoff distance in Angstroms, by default 6.0
    """
    self.cutoff = cutoff
calculate(atoms, ethresh=1e-06)

Calculate Coulomb energy using PME.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object with positions and charges

 required
ethresh float

Ewald threshold, by default 1e-6

 1e-06

Returns:

Type Description
float

Coulomb energy in eV/particle
Source code in toolbox/calculator/coulomb.py 
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
300
301
302
303
def calculate(self, atoms, ethresh=1e-6) -> float:
    """Calculate Coulomb energy using PME.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object with positions and charges
    ethresh : float, optional
        Ewald threshold, by default 1e-6

    Returns
    -------
    float
        Coulomb energy in eV/particle
    """
    positions = atoms.get_positions()
    box = atoms.get_cell()
    charges = atoms.get_initial_charges()
    pairs = calculate_pairs(atoms, self.cutoff)

    box = torch.tensor(box, dtype=torch.double, requires_grad=False, device=DEVICE)
    positions = torch.tensor(
        positions, dtype=torch.double, requires_grad=True, device=DEVICE
    )
    charges = torch.tensor(
        charges, dtype=torch.double, requires_grad=False, device=DEVICE
    )
    pairs = torch.tensor(
        pairs, dtype=torch.int32, requires_grad=False, device=DEVICE
    )

    mscales = torch.tensor([0.0, 0.0, 1.0, 1.0, 1.0, 1.0], device=DEVICE)
    kappa, K = self.setup_ewald(box, ethresh)
    K1, K2, K3 = K
    pme_recip_fn = generate_pme_recip(
        Ck_fn=Ck_1,
        kappa=kappa,
        gamma=False,
        pme_order=6,
        K1=K1,
        K2=K2,
        K3=K3,
        lmax=0,
    )
    charges = torch.reshape(charges, (-1, 1))
    energy = energy_pme(
        positions,
        box,
        pairs,
        charges,
        None,
        None,
        None,
        mscales,
        None,
        None,
        None,
        pme_recip_fn,
        kappa,
        K1,
        K2,
        K3,
        0,
        False,
        True,
    )[0]
    # from kJ/mol to eV/particle
    return (energy * ENERGY_COEFF).item()
setup_ewald(box, ethresh)

Set up Ewald parameters for PME calculation.

Parameters:

Name Type Description Default
box array_like

Simulation box vectors

 required
ethresh float

Ewald threshold

 required

Returns:

Type Description
tuple

Tuple of (kappa, K) where kappa is the screening parameter and K is the grid dimensions
Source code in toolbox/calculator/coulomb.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
def setup_ewald(self, box, ethresh):
    """Set up Ewald parameters for PME calculation.

    Parameters
    ----------
    box : array_like
        Simulation box vectors
    ethresh : float
        Ewald threshold

    Returns
    -------
    tuple
        Tuple of (kappa, K) where kappa is the screening parameter
        and K is the grid dimensions
    """
    kappa, K1, K2, K3 = setup_ewald_parameters(
        torch.tensor(self.cutoff), torch.tensor(ethresh), box, 0.01, "openmm"
    )
    K = (K1, K2, K3)
    return kappa, K

calculate_pairs(atoms, rcut)

Calculate atom pairs within cutoff distance.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object

 required
rcut float

Cutoff distance in Angstroms

 required

Returns:

Type Description
list

List of atom pairs within cutoff, each as [i, j, 0]
Source code in toolbox/calculator/coulomb.py 
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
def calculate_pairs(atoms, rcut):
    """Calculate atom pairs within cutoff distance.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object
    rcut : float
        Cutoff distance in Angstroms

    Returns
    -------
    list
        List of atom pairs within cutoff, each as [i, j, 0]
    """
    cellpar = atoms.cell.cellpar()
    positions = atoms.get_positions()

    pairs = []
    dist_mat = distance_array(positions, positions, box=cellpar)
    for ii, dist_vec in enumerate(dist_mat):
        mask = dist_vec < rcut
        jjs = np.where(mask)[0]
        mask = jjs > ii
        jjs = jjs[mask]
        for jj in jjs:
            pairs.append([ii, jj, 0])
    return pairs

coul(qi, qj, rij)

Calculate Coulomb energy between two point charges.

Parameters:

Name Type Description Default
qi float

First charge in elementary charge units (e)

 required
qj float

Second charge in elementary charge units (e)

 required
rij float

Distance between charges in Angstroms (A)

 required

Returns:

Type Description
float

Coulomb energy in electron volts (eV)
Source code in toolbox/calculator/coulomb.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def coul(qi, qj, rij):
    """Calculate Coulomb energy between two point charges.

    Parameters
    ----------
    qi : float
        First charge in elementary charge units (e)
    qj : float
        Second charge in elementary charge units (e)
    rij : float
        Distance between charges in Angstroms (A)

    Returns
    -------
    float
        Coulomb energy in electron volts (eV)
    """
    e = qqrd2e * qi * qj / rij
    return e

real_coul(qi, qj, xi, xj, gewald, force=False)

Calculate real-space Coulomb energy with Ewald summation.

Parameters:

Name Type Description Default
qi float

First charge in elementary charge units (e)

 required
qj float

Second charge in elementary charge units (e)

 required
xi array_like

Position of first charge

 required
xj array_like

Position of second charge

 required
gewald float

Ewald screening parameter

 required
force bool

Whether to calculate forces, by default False

 False

Returns:

Type Description
float or tuple

If force=False, returns Coulomb energy in eV If force=True, returns tuple of (energy, force)
Source code in toolbox/calculator/coulomb.py 
 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
def real_coul(qi, qj, xi, xj, gewald, force=False):
    """Calculate real-space Coulomb energy with Ewald summation.

    Parameters
    ----------
    qi : float
        First charge in elementary charge units (e)
    qj : float
        Second charge in elementary charge units (e)
    xi : array_like
        Position of first charge
    xj : array_like
        Position of second charge
    gewald : float
        Ewald screening parameter
    force : bool, optional
        Whether to calculate forces, by default False

    Returns
    -------
    float or tuple
        If force=False, returns Coulomb energy in eV
        If force=True, returns tuple of (energy, force)
    """
    rij = np.linalg.norm(xi - xj)
    prefactor = qqrd2e * qi * qj / rij
    if force:
        grij = gewald * rij
        expm2 = np.exp(-grij * grij)
        t = 1.0 / (1.0 + EWALD_P * grij)
        erfc = t * (A1 + t * (A2 + t * (A3 + t * (A4 + t * A5)))) * expm2
        forcecoul = prefactor * (erfc + EWALD_F * grij * expm2)
        fpair = forcecoul / rij**2
        f = (xi - xj) * fpair
        e = prefactor * erfc
        return e, f
    else:
        e = prefactor * special.erfc(rij * gewald)
        return e

CP2K Calculator

toolbox.calculator.cp2k

CP2K quantum chemistry calculator interface.

This module provides a calculator interface for running CP2K quantum chemistry calculations with proper error handling and output checking.

Cp2kCalculator

Bases: BashCalculator

A calculator for running CP2K quantum chemistry calculations.

This class extends BashCalculator to provide specific functionality for CP2K calculations, including error handling and output checking.

Source code in toolbox/calculator/cp2k.py 
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
class Cp2kCalculator(BashCalculator):
    """A calculator for running CP2K quantum chemistry calculations.

    This class extends BashCalculator to provide specific functionality for
    CP2K calculations, including error handling and output checking.
    """

    def __init__(self, work_dir) -> None:
        """Initialize the Cp2kCalculator.

        Parameters
        ----------
        work_dir : str
            The working directory where CP2K calculations will be executed
        """
        super().__init__(work_dir)

    def run(
        self,
        command="cp2k.popt",
        stdin="input.inp",
        stdout="output.out",
        stderr="cp2k.stderr",
        mpi_command: str = "mpiexec.hydra",
        ignore_finished_tag: bool = False,
        ignore_err: bool = True,
    ):
        """Run a CP2K calculation.

        Parameters
        ----------
        command : str, optional
            CP2K executable command, by default "cp2k.popt"
        stdin : str, optional
            Input file name, by default "input.inp"
        stdout : str, optional
            Standard output file name, by default "output.out"
        stderr : str, optional
            Standard error file name, by default "cp2k.stderr"
        mpi_command : str, optional
            MPI command to use, by default "mpiexec.hydra"
        ignore_finished_tag : bool, optional
            Whether to ignore existing finished tag, by default False
        ignore_err : bool, optional
            Whether to ignore calculation errors, by default True
        """
        self.ignore_err = ignore_err
        super().run(command, stdin, stdout, stderr, mpi_command, ignore_finished_tag)

    def _make_finished_tag(self, stdout):
        """Create a finished tag after successful CP2K calculation.

        Parameters
        ----------
        stdout : str
            Name of the CP2K output file to check for completion

        Raises
        ------
        SystemExit
            If CP2K calculation did not finish and ignore_err is False
        """
        try:
            Cp2kOutput(stdout)
            open(os.path.join("finished_tag"), "w").close()
        except Exception:
            warning_msg = "CP2K calculation does not finish!"
            if self.ignore_err:
                logging.warning(warning_msg)
            else:
                sys.exit(warning_msg)
__init__(work_dir)

Initialize the Cp2kCalculator.

Parameters:

Name Type Description Default
work_dir str

The working directory where CP2K calculations will be executed

 required
Source code in toolbox/calculator/cp2k.py 
23
24
25
26
27
28
29
30
31
def __init__(self, work_dir) -> None:
    """Initialize the Cp2kCalculator.

    Parameters
    ----------
    work_dir : str
        The working directory where CP2K calculations will be executed
    """
    super().__init__(work_dir)
run(command='cp2k.popt', stdin='input.inp', stdout='output.out', stderr='cp2k.stderr', mpi_command='mpiexec.hydra', ignore_finished_tag=False, ignore_err=True)

Run a CP2K calculation.

Parameters:

Name Type Description Default
command str

CP2K executable command, by default "cp2k.popt"

 'cp2k.popt'
stdin str

Input file name, by default "input.inp"

 'input.inp'
stdout str

Standard output file name, by default "output.out"

 'output.out'
stderr str

Standard error file name, by default "cp2k.stderr"

 'cp2k.stderr'
mpi_command str

MPI command to use, by default "mpiexec.hydra"

 'mpiexec.hydra'
ignore_finished_tag bool

Whether to ignore existing finished tag, by default False

 False
ignore_err bool

Whether to ignore calculation errors, by default True

 True
Source code in toolbox/calculator/cp2k.py 
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
def run(
    self,
    command="cp2k.popt",
    stdin="input.inp",
    stdout="output.out",
    stderr="cp2k.stderr",
    mpi_command: str = "mpiexec.hydra",
    ignore_finished_tag: bool = False,
    ignore_err: bool = True,
):
    """Run a CP2K calculation.

    Parameters
    ----------
    command : str, optional
        CP2K executable command, by default "cp2k.popt"
    stdin : str, optional
        Input file name, by default "input.inp"
    stdout : str, optional
        Standard output file name, by default "output.out"
    stderr : str, optional
        Standard error file name, by default "cp2k.stderr"
    mpi_command : str, optional
        MPI command to use, by default "mpiexec.hydra"
    ignore_finished_tag : bool, optional
        Whether to ignore existing finished tag, by default False
    ignore_err : bool, optional
        Whether to ignore calculation errors, by default True
    """
    self.ignore_err = ignore_err
    super().run(command, stdin, stdout, stderr, mpi_command, ignore_finished_tag)

D3 Calculator

toolbox.calculator.d3

D3 dispersion correction calculator.

This module provides a calculator for Grimme's D3 dispersion corrections using the tad_dftd3 library.

D3Calculator

Calculator for D3 dispersion corrections.

This calculator computes Grimme's D3 dispersion corrections using the tad_dftd3 library.

Source code in toolbox/calculator/d3.py 
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
66
67
68
class D3Calculator:
    """Calculator for D3 dispersion corrections.

    This calculator computes Grimme's D3 dispersion corrections
    using the tad_dftd3 library.
    """

    def __init__(self, params: Optional[dict] = None) -> None:
        """Initialize D3Calculator.

        Parameters
        ----------
        params : Optional[Dict], optional
            D3 parameters dictionary, by default None.
            If None, uses r²SCAN-D3(BJ) parameters.
        """
        self.ref = d3.reference.Reference()
        if params is None:
            # r²SCAN-D3(BJ)
            self.params = dict(
                a1=torch.tensor(0.49484001),
                s8=torch.tensor(0.78981345),
                a2=torch.tensor(5.73083694),
            )
        else:
            self.params = params

    def run(self, atoms: Atoms):
        """Calculate D3 dispersion energy.

        Parameters
        ----------
        atoms : ase.Atoms
            ASE Atoms object with atomic numbers and positions

        Returns
        -------
        torch.Tensor
            D3 dispersion energy
        """
        # nframes * natoms
        numbers = atoms.get_atomic_numbers().reshape(-1, len(atoms))
        numbers = torch.tensor(numbers, dtype=torch.int64)
        # nframes * 3 * natoms
        positions = atoms.get_positions().reshape(-1, len(atoms), 3)
        positions = torch.tensor(positions)

        cn = d3.ncoord.cn_d3(numbers, positions)
        weights = d3.model.weight_references(numbers, cn, self.ref)
        c6 = d3.model.atomic_c6(numbers, weights, self.ref)
        energy = d3.disp.dispersion(numbers, positions, self.params, c6)
        return energy.sum()
__init__(params=None)

Initialize D3Calculator.

Parameters:

Name Type Description Default
params Optional[Dict]

D3 parameters dictionary, by default None. If None, uses r²SCAN-D3(BJ) parameters.

 None
Source code in toolbox/calculator/d3.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(self, params: Optional[dict] = None) -> None:
    """Initialize D3Calculator.

    Parameters
    ----------
    params : Optional[Dict], optional
        D3 parameters dictionary, by default None.
        If None, uses r²SCAN-D3(BJ) parameters.
    """
    self.ref = d3.reference.Reference()
    if params is None:
        # r²SCAN-D3(BJ)
        self.params = dict(
            a1=torch.tensor(0.49484001),
            s8=torch.tensor(0.78981345),
            a2=torch.tensor(5.73083694),
        )
    else:
        self.params = params
run(atoms)

Calculate D3 dispersion energy.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object with atomic numbers and positions

 required

Returns:

Type Description
Tensor

D3 dispersion energy
Source code in toolbox/calculator/d3.py 
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
def run(self, atoms: Atoms):
    """Calculate D3 dispersion energy.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object with atomic numbers and positions

    Returns
    -------
    torch.Tensor
        D3 dispersion energy
    """
    # nframes * natoms
    numbers = atoms.get_atomic_numbers().reshape(-1, len(atoms))
    numbers = torch.tensor(numbers, dtype=torch.int64)
    # nframes * 3 * natoms
    positions = atoms.get_positions().reshape(-1, len(atoms), 3)
    positions = torch.tensor(positions)

    cn = d3.ncoord.cn_d3(numbers, positions)
    weights = d3.model.weight_references(numbers, cn, self.ref)
    c6 = d3.model.atomic_c6(numbers, weights, self.ref)
    energy = d3.disp.dispersion(numbers, positions, self.params, c6)
    return energy.sum()

Deep Potential Calculator

toolbox.calculator.dp

Deep Potential dispatcher module.

This module provides classes for dispatching Deep Potential calculations to HPC systems using dpdispatcher.

CP2KDPDispatcher

Bases: DPDispatcher

Dispatcher for running CP2K calculations with Deep Potential.

This class extends DPDispatcher to provide specific functionality for CP2K calculations with version-specific configurations.

Source code in toolbox/calculator/dp.py 
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
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
class CP2KDPDispatcher(DPDispatcher):
    """Dispatcher for running CP2K calculations with Deep Potential.

    This class extends DPDispatcher to provide specific functionality
    for CP2K calculations with version-specific configurations.
    """

    def __init__(self, work_dir=".", v9: bool = False) -> None:
        """Initialize CP2KDPDispatcher.

        Parameters
        ----------
        work_dir : str, optional
            Working directory for calculations, by default "."
        v9 : bool, optional
            Whether to use CP2K v9 configuration, by default False
        """
        super().__init__(work_dir)
        self.v9 = v9

    def _setup(self, machine_setup=None, resources_setup=None):
        """Set up CP2K-specific machine and resources configuration.

        Parameters
        ----------
        machine_setup : dict, optional
            Machine configuration overrides, by default {}
        resources_setup : dict, optional
            Resources configuration overrides, by default {}
        """
        if resources_setup is None:
            resources_setup = {}
        if machine_setup is None:
            machine_setup = {}
        if self.v9:
            _resources_setup = {
                "custom_flags": ["#SBATCH -J cp2k"],
                "module_list": ["mkl/latest", "mpi/latest", "gcc/9.3.0", "cp2k/9.1"],
            }
        else:
            _resources_setup = {
                "custom_flags": ["#SBATCH -J cp2k"],
                "module_list": [
                    "mpi/intel/2017.5.239",
                    "intel/17.5.239",
                    "gcc/7.4.0",
                    "cp2k/7.1",
                ],
            }
        _resources_setup.update(resources_setup)
        super()._setup(machine_setup, _resources_setup)

    def run(self, dnames, machine_setup=None, resources_setup=None, task_setup=None):
        """Run CP2K calculations on specified directories.

        Parameters
        ----------
        dnames : list
            List of directory names to run calculations in
        machine_setup : dict, optional
            Machine configuration overrides, by default {}
        resources_setup : dict, optional
            Resources configuration overrides, by default {}
        task_setup : dict, optional
            Task configuration overrides, by default {}
        """
        if task_setup is None:
            task_setup = {}
        if resources_setup is None:
            resources_setup = {}
        if machine_setup is None:
            machine_setup = {}
        _task_setup = {
            "command": "mpiexec.hydra cp2k.popt input.inp",
            "forward_files": ["input.inp", "coord.xyz"],
            "backward_files": [
                "output",
                "cp2k-v_hartree-1_0.cube",
                "cp2k-ELECTRON_DENSITY-1_0.cube",
                "cp2k-TOTAL_DENSITY-1_0.cube",
                "cp2k-RESTART.wfn",
            ],
            "outlog": "output",
        }
        _task_setup.update(task_setup)
        super().run(dnames, machine_setup, resources_setup, _task_setup)
__init__(work_dir='.', v9=False)

Initialize CP2KDPDispatcher.

Parameters:

Name Type Description Default
work_dir str

Working directory for calculations, by default "."

 '.'
v9 bool

Whether to use CP2K v9 configuration, by default False

 False
Source code in toolbox/calculator/dp.py 
159
160
161
162
163
164
165
166
167
168
169
170
def __init__(self, work_dir=".", v9: bool = False) -> None:
    """Initialize CP2KDPDispatcher.

    Parameters
    ----------
    work_dir : str, optional
        Working directory for calculations, by default "."
    v9 : bool, optional
        Whether to use CP2K v9 configuration, by default False
    """
    super().__init__(work_dir)
    self.v9 = v9
run(dnames, machine_setup=None, resources_setup=None, task_setup=None)

Run CP2K calculations on specified directories.

Parameters:

Name Type Description Default
dnames list

List of directory names to run calculations in

 required
machine_setup dict

Machine configuration overrides, by default {}

 None
resources_setup dict

Resources configuration overrides, by default {}

 None
task_setup dict

Task configuration overrides, by default {}

 None
Source code in toolbox/calculator/dp.py 
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
def run(self, dnames, machine_setup=None, resources_setup=None, task_setup=None):
    """Run CP2K calculations on specified directories.

    Parameters
    ----------
    dnames : list
        List of directory names to run calculations in
    machine_setup : dict, optional
        Machine configuration overrides, by default {}
    resources_setup : dict, optional
        Resources configuration overrides, by default {}
    task_setup : dict, optional
        Task configuration overrides, by default {}
    """
    if task_setup is None:
        task_setup = {}
    if resources_setup is None:
        resources_setup = {}
    if machine_setup is None:
        machine_setup = {}
    _task_setup = {
        "command": "mpiexec.hydra cp2k.popt input.inp",
        "forward_files": ["input.inp", "coord.xyz"],
        "backward_files": [
            "output",
            "cp2k-v_hartree-1_0.cube",
            "cp2k-ELECTRON_DENSITY-1_0.cube",
            "cp2k-TOTAL_DENSITY-1_0.cube",
            "cp2k-RESTART.wfn",
        ],
        "outlog": "output",
    }
    _task_setup.update(task_setup)
    super().run(dnames, machine_setup, resources_setup, _task_setup)

DPDispatcher

Dispatcher for running Deep Potential calculations.

This class provides functionality to dispatch Deep Potential calculations to HPC systems using dpdispatcher.

Source code in toolbox/calculator/dp.py 
 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
 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class DPDispatcher:
    """Dispatcher for running Deep Potential calculations.

    This class provides functionality to dispatch Deep Potential
    calculations to HPC systems using dpdispatcher.
    """

    def __init__(self, work_dir=".") -> None:
        """Initialize DPDispatcher.

        Parameters
        ----------
        work_dir : str, optional
            Working directory for calculations, by default "."
        """
        self.work_dir = work_dir

    def _setup(self, machine_setup=None, resources_setup=None):
        """Set up machine and resources configuration.

        Parameters
        ----------
        machine_setup : dict, optional
            Machine configuration overrides, by default {}
        resources_setup : dict, optional
            Resources configuration overrides, by default {}
        """
        if resources_setup is None:
            resources_setup = {}
        if machine_setup is None:
            machine_setup = {}
        _machine_setup = {
            "batch_type": "Slurm",
            "context_type": "LocalContext",
            "local_root": "./",
            "remote_root": "/data/jxzhu/nnp/dp_workdir",
        }
        _machine_setup.update(machine_setup)
        _resources_setup = {
            "number_node": 1,
            "cpu_per_node": 24,
            "gpu_per_node": 0,
            "kwargs": {"gpu_usage": False},
            "queue_name": "c51-large",
            "group_size": 80,
            "module_purge": True,
        }
        _resources_setup.update(resources_setup)
        self._check_node(_resources_setup)

        self.machine = Machine(**_machine_setup)
        self.resources = Resources(**_resources_setup)

    def run(self, dnames, machine_setup=None, resources_setup=None, task_setup=None):
        """Run calculations on specified directories.

        Parameters
        ----------
        dnames : list
            List of directory names to run calculations in
        machine_setup : dict, optional
            Machine configuration overrides, by default {}
        resources_setup : dict, optional
            Resources configuration overrides, by default {}
        task_setup : dict, optional
            Task configuration overrides, by default {}
        """
        if task_setup is None:
            task_setup = {}
        if resources_setup is None:
            resources_setup = {}
        if machine_setup is None:
            machine_setup = {}
        self._setup(machine_setup, resources_setup)

        task_list = []
        for dname in dnames:
            task = Task(task_work_path=dname, **task_setup)
            task_list.append(task)

        submission = Submission(
            work_base=self.work_dir,
            machine=self.machine,
            resources=self.resources,
            task_list=task_list,
        )
        submission.run_submission()

    @staticmethod
    def _check_node(resources_setup):
        """Check and adjust resources based on queue name.

        Parameters
        ----------
        resources_setup : dict
            Resources configuration to check and modify

        Raises
        ------
        AttributeError
            If queue name is unknown
        """
        queue = resources_setup["queue_name"]

        if "c51" in queue:
            resources_setup["cpu_per_node"] = 24
        elif "c52" in queue:
            resources_setup["cpu_per_node"] = 28
        elif "c53" in queue:
            resources_setup["cpu_per_node"] = 32
        elif queue == "cpu":
            # add queue in ikkem
            cpu_per_node = resources_setup.get("cpu_per_node", 64)
            resources_setup["cpu_per_node"] = cpu_per_node
        else:
            raise AttributeError(f"Unknown queue: {queue}")
__init__(work_dir='.')

Initialize DPDispatcher.

Parameters:

Name Type Description Default
work_dir str

Working directory for calculations, by default "."

 '.'
Source code in toolbox/calculator/dp.py 
18
19
20
21
22
23
24
25
26
def __init__(self, work_dir=".") -> None:
    """Initialize DPDispatcher.

    Parameters
    ----------
    work_dir : str, optional
        Working directory for calculations, by default "."
    """
    self.work_dir = work_dir
run(dnames, machine_setup=None, resources_setup=None, task_setup=None)

Run calculations on specified directories.

Parameters:

Name Type Description Default
dnames list

List of directory names to run calculations in

 required
machine_setup dict

Machine configuration overrides, by default {}

 None
resources_setup dict

Resources configuration overrides, by default {}

 None
task_setup dict

Task configuration overrides, by default {}

 None
Source code in toolbox/calculator/dp.py 
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
def run(self, dnames, machine_setup=None, resources_setup=None, task_setup=None):
    """Run calculations on specified directories.

    Parameters
    ----------
    dnames : list
        List of directory names to run calculations in
    machine_setup : dict, optional
        Machine configuration overrides, by default {}
    resources_setup : dict, optional
        Resources configuration overrides, by default {}
    task_setup : dict, optional
        Task configuration overrides, by default {}
    """
    if task_setup is None:
        task_setup = {}
    if resources_setup is None:
        resources_setup = {}
    if machine_setup is None:
        machine_setup = {}
    self._setup(machine_setup, resources_setup)

    task_list = []
    for dname in dnames:
        task = Task(task_work_path=dname, **task_setup)
        task_list.append(task)

    submission = Submission(
        work_base=self.work_dir,
        machine=self.machine,
        resources=self.resources,
        task_list=task_list,
    )
    submission.run_submission()

Electric Potential Calculator

toolbox.calculator.elecpot

Electrostatic potential calculator module.

This module provides classes for calculating electrostatic potential from charge density using various boundary conditions.

Reference: - https://altafang.com/2020/10/13/calculating-1d-electrostatic-potential-profile-from-md-simulations/ - Code from Justina Moss (Leiden University, Email: j.h.moss@lic.leidenuniv.nl) - https://github.com/SINGROUP/Potential_solver.

ElecPotentialCalculator

Calculator for 1D electrostatic potential from charge density.

This calculator computes electrostatic potential from charge density using various boundary conditions.

Parameters:

Name Type Description Default
charge array_like

Charge density in e/ų

 required
grid array_like

Grid points in Ã…

 required

Returns:

Type Description
Electrostatic potential in eV

(different from Hartree potential with a negative sign!)
Source code in toolbox/calculator/elecpot.py 
 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
 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
112
113
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
class ElecPotentialCalculator:
    """Calculator for 1D electrostatic potential from charge density.

    This calculator computes electrostatic potential from charge density
    using various boundary conditions.

    Parameters
    ----------
    charge : array_like
        Charge density in e/ų
    grid : array_like
        Grid points in Ã…

    Returns
    -------
    Electrostatic potential in eV
        (different from Hartree potential with a negative sign!)
    """

    def __init__(self, charge, grid) -> None:
        """Initialize ElecPotentialCalculator.

        Parameters
        ----------
        charge : array_like
            Charge density in e/ų
        grid : array_like
            Grid points in Ã…
        """
        if len(grid) != len(charge):
            raise AttributeError("Grid and charge should have the same dimension.")
        self.grid = grid
        self.charge = charge

    def calculate(self, bc="periodic", **kwargs):
        """Calculate electrostatic potential.

        Parameters
        ----------
        bc : str, optional
            Boundary condition, by default "periodic"
            Options: "periodic", "open", "dirichlet", "neumann", "dip_cor"
        **kwargs
            Additional keyword arguments for specific boundary conditions

        Returns
        -------
        array_like
            Electrostatic potential in V

        Raises
        ------
        AttributeError
            If boundary condition is not supported
        """
        if (not hasattr(self, "int1")) or (not hasattr(self, "int2")):
            self._integrate()
        self.bc = bc
        try:
            self.potential = getattr(self, f"_calculate_{bc}")(**kwargs)
            return self.potential
        except AttributeError as e:
            raise AttributeError(f"Unsupported boundary condition {bc}") from e

    def _integrate(self):
        """Perform double integration of charge density."""
        self.int1 = integrate.cumulative_trapezoid(self.charge, self.grid, initial=0)
        self.int2 = (
            -integrate.cumulative_trapezoid(self.int1, self.grid, initial=0) / EPSILON
        )

    def _calculate_periodic(self, l_box):
        """Calculate potential with periodic boundary conditions.

        Parameters
        ----------
        l_box : float
            Box length in Ã…
        """
        phi = self.solve_fft_poisson(l_box)
        return phi

    def _calculate_open(self):
        """Calculate potential with open boundary conditions."""
        return self.int2

    def _calculate_dirichlet(self):
        """Calculate potential with Dirichlet boundary conditions.

        Raises
        ------
        NotImplementedError
            This method is not implemented
        """
        raise NotImplementedError

    def _calculate_neumann(self):
        """Calculate potential with Neumann boundary conditions.

        Raises
        ------
        NotImplementedError
            This method is not implemented
        """
        raise NotImplementedError

    def _calculate_dip_cor(self, cell):
        """Calculate potential with dipole correction.

        Parameters
        ----------
        cell : array_like
            Unit cell vectors or parameters

        Returns
        -------
        array_like
            Potential with dipole correction applied
        """
        if np.shape(cell) == (3,):
            cell = np.diag(cell)
        elif np.shape(cell) == (6,):
            cell = cellpar_to_cell(cell)
        elif np.shape(cell) == (3, 3):
            pass
        else:
            raise AttributeError("")

        z_max = cell[2][2]
        phi = self._calculate_periodic(z_max)
        cross_area = np.linalg.norm(np.cross(cell[0], cell[1]))
        surf_pol = np.sum(self.grid * self.charge) / cross_area
        v_cor = surf_pol / EPSILON * (self.grid / z_max - 0.5)
        return phi + v_cor

    def solve_fft_poisson(self, l_box):
        """Solve Poisson equation using FFT.

        Parameters
        ----------
        l_box : float
            Box length in Ã…

        Returns
        -------
        array_like
            Electrostatic potential
        """
        n_grid = len(self.grid)
        grid_spacing = l_box / n_grid

        start = time.process_time()
        charge_k_space = np.fft.fft(self.charge)
        end = time.process_time()
        print(f"| | FFT took {end - start} s")

        start = time.process_time()
        pot_k_space = np.zeros(n_grid, dtype=complex)
        k = np.fft.fftfreq(n_grid, grid_spacing)
        k_squared = k**2
        pot_k_space[1:] = charge_k_space[1:] / k_squared[1:]
        pot_k_space = pot_k_space / (4.0 * np.pi * np.pi * EPSILON)
        end = time.process_time()
        print(f"| | k-space arithmetics took {end - start} s")

        start = time.process_time()
        pot_grid = np.fft.ifftn(pot_k_space).real
        end = time.process_time()
        print(f"| | Inverse FFT took {end - start} s")
        return pot_grid
__init__(charge, grid)

Initialize ElecPotentialCalculator.

Parameters:

Name Type Description Default
charge array_like

Charge density in e/ų

 required
grid array_like

Grid points in Ã…

 required
Source code in toolbox/calculator/elecpot.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(self, charge, grid) -> None:
    """Initialize ElecPotentialCalculator.

    Parameters
    ----------
    charge : array_like
        Charge density in e/ų
    grid : array_like
        Grid points in Ã…
    """
    if len(grid) != len(charge):
        raise AttributeError("Grid and charge should have the same dimension.")
    self.grid = grid
    self.charge = charge
calculate(bc='periodic', **kwargs)

Calculate electrostatic potential.

Parameters:

Name Type Description Default
bc str

Boundary condition, by default "periodic" Options: "periodic", "open", "dirichlet", "neumann", "dip_cor"

 'periodic'
**kwargs

Additional keyword arguments for specific boundary conditions

 {}

Returns:

Type Description
array_like

Electrostatic potential in V

Raises:

Type Description
AttributeError

If boundary condition is not supported
Source code in toolbox/calculator/elecpot.py 
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
def calculate(self, bc="periodic", **kwargs):
    """Calculate electrostatic potential.

    Parameters
    ----------
    bc : str, optional
        Boundary condition, by default "periodic"
        Options: "periodic", "open", "dirichlet", "neumann", "dip_cor"
    **kwargs
        Additional keyword arguments for specific boundary conditions

    Returns
    -------
    array_like
        Electrostatic potential in V

    Raises
    ------
    AttributeError
        If boundary condition is not supported
    """
    if (not hasattr(self, "int1")) or (not hasattr(self, "int2")):
        self._integrate()
    self.bc = bc
    try:
        self.potential = getattr(self, f"_calculate_{bc}")(**kwargs)
        return self.potential
    except AttributeError as e:
        raise AttributeError(f"Unsupported boundary condition {bc}") from e
solve_fft_poisson(l_box)

Solve Poisson equation using FFT.

Parameters:

Name Type Description Default
l_box float

Box length in Ã…

 required

Returns:

Type Description
array_like

Electrostatic potential
Source code in toolbox/calculator/elecpot.py 
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
def solve_fft_poisson(self, l_box):
    """Solve Poisson equation using FFT.

    Parameters
    ----------
    l_box : float
        Box length in Ã…

    Returns
    -------
    array_like
        Electrostatic potential
    """
    n_grid = len(self.grid)
    grid_spacing = l_box / n_grid

    start = time.process_time()
    charge_k_space = np.fft.fft(self.charge)
    end = time.process_time()
    print(f"| | FFT took {end - start} s")

    start = time.process_time()
    pot_k_space = np.zeros(n_grid, dtype=complex)
    k = np.fft.fftfreq(n_grid, grid_spacing)
    k_squared = k**2
    pot_k_space[1:] = charge_k_space[1:] / k_squared[1:]
    pot_k_space = pot_k_space / (4.0 * np.pi * np.pi * EPSILON)
    end = time.process_time()
    print(f"| | k-space arithmetics took {end - start} s")

    start = time.process_time()
    pot_grid = np.fft.ifftn(pot_k_space).real
    end = time.process_time()
    print(f"| | Inverse FFT took {end - start} s")
    return pot_grid

GaussianElecPotentialCalculator

Calculator for electrostatic potential from Gaussian charge distributions.

This calculator computes electrostatic potential from atoms represented as Gaussian charge distributions.

Source code in toolbox/calculator/elecpot.py 
198
199
200
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
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
300
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
class GaussianElecPotentialCalculator:
    """Calculator for electrostatic potential from Gaussian charge distributions.

    This calculator computes electrostatic potential from atoms
    represented as Gaussian charge distributions.
    """

    def __init__(
        self,
        grids: np.ndarray,
        l_box: float,
        cross_area: float,
        spread_dict: Optional[dict[str, float]] = None,
        charge_dict: Optional[dict[str, float]] = None,
    ) -> None:
        """Initialize GaussianElecPotentialCalculator.

        Parameters
        ----------
        grids : np.ndarray
            Grid points in Ã…
        l_box : float
            Box length in Ã…
        cross_area : float
            Cross-sectional area in Ų
        spread_dict : Optional[Dict[str, float]], optional
            Dictionary of Gaussian spreads for each element, by default None
        charge_dict : Optional[Dict[str, float]], optional
            Dictionary of charges for each element, by default None
        """
        # setup grid
        self.l_box = l_box
        self.grids = grids
        self.grid_edges = np.linspace(0.0, self.l_box, len(self.grids) + 1)
        dx = np.diff(self.grid_edges)[0]
        self.grid_volume = cross_area * dx

        self.spread_dict = spread_dict
        self.charge_dict = charge_dict

    def calc_rho(
        self,
        mu: np.ndarray,
        spread: np.ndarray,
        charge: np.ndarray,
    ):
        """Calculate charge density from Gaussian distributions.

        Parameters
        ----------
        mu : np.ndarray
            Centers of Gaussian distributions
        spread : np.ndarray
            Spreads of Gaussian distributions
        charge : np.ndarray
            Charges of Gaussian distributions

        Returns
        -------
        np.ndarray
            Charge density on the grid
        """
        grid_edges = np.reshape(self.grid_edges, (1, -1))
        spread = np.reshape(spread, [-1, 1])
        mu = np.reshape(mu, [-1, 1])
        charge = np.reshape(charge, [-1, 1])

        # nat * ngrid
        out = gaussian_int(grid_edges[:, 1:], mu, spread) - gaussian_int(
            grid_edges[:, :-1], mu, spread
        )
        out = np.sum(out * charge, axis=0)
        # deal with periodic boundary
        # left image
        l_out = gaussian_int(grid_edges[:, 1:] - self.l_box, mu, spread) - gaussian_int(
            grid_edges[:, :-1] - self.l_box, mu, spread
        )
        l_out = np.sum(l_out * charge, axis=0)
        out += l_out
        # right image
        r_out = gaussian_int(grid_edges[:, 1:] + self.l_box, mu, spread) - gaussian_int(
            grid_edges[:, :-1] + self.l_box, mu, spread
        )
        r_out = np.sum(r_out * charge, axis=0)
        out += r_out

        np.testing.assert_almost_equal(out.sum(), charge.sum())

        rho = out / self.grid_volume
        return rho

    def run(
        self,
        mu: Optional[np.ndarray] = None,
        spread: Optional[np.ndarray] = None,
        charge: Optional[np.ndarray] = None,
        atoms: Optional[Atoms] = None,
    ):
        """Calculate electrostatic potential from Gaussian charge distributions.

        Parameters
        ----------
        mu : Optional[np.ndarray], optional
            Centers of Gaussian distributions, by default None
        spread : Optional[np.ndarray], optional
            Spreads of Gaussian distributions, by default None
        charge : Optional[np.ndarray], optional
            Charges of Gaussian distributions, by default None
        atoms : Optional[Atoms], optional
            ASE Atoms object, by default None

        Returns
        -------
        np.ndarray
            Electrostatic potential
        """
        if mu is None:
            mu = atoms.get_positions()[:, 2]
        if spread is None:
            assert self.spread_dict is not None, "spread_dict is not set"
            spread = np.array([self.spread_dict[s] for s in atoms.symbols])
        if charge is None:
            if self.charge_dict is not None:
                charge = np.array([self.charge_dict[s] for s in atoms.symbols])
            else:
                charge = atoms.get_initial_charges()

        rho = self.calc_rho(mu, spread, charge)
        calculator = ElecPotentialCalculator(rho, self.grids)
        phi = calculator.calculate(l_box=self.l_box)
        return phi
__init__(grids, l_box, cross_area, spread_dict=None, charge_dict=None)

Initialize GaussianElecPotentialCalculator.

Parameters:

Name Type Description Default
grids ndarray

Grid points in Ã…

 required
l_box float

Box length in Ã…

 required
cross_area float

Cross-sectional area in Ų

 required
spread_dict Optional[Dict[str, float]]

Dictionary of Gaussian spreads for each element, by default None

 None
charge_dict Optional[Dict[str, float]]

Dictionary of charges for each element, by default None

 None
Source code in toolbox/calculator/elecpot.py 
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
def __init__(
    self,
    grids: np.ndarray,
    l_box: float,
    cross_area: float,
    spread_dict: Optional[dict[str, float]] = None,
    charge_dict: Optional[dict[str, float]] = None,
) -> None:
    """Initialize GaussianElecPotentialCalculator.

    Parameters
    ----------
    grids : np.ndarray
        Grid points in Ã…
    l_box : float
        Box length in Ã…
    cross_area : float
        Cross-sectional area in Ų
    spread_dict : Optional[Dict[str, float]], optional
        Dictionary of Gaussian spreads for each element, by default None
    charge_dict : Optional[Dict[str, float]], optional
        Dictionary of charges for each element, by default None
    """
    # setup grid
    self.l_box = l_box
    self.grids = grids
    self.grid_edges = np.linspace(0.0, self.l_box, len(self.grids) + 1)
    dx = np.diff(self.grid_edges)[0]
    self.grid_volume = cross_area * dx

    self.spread_dict = spread_dict
    self.charge_dict = charge_dict
calc_rho(mu, spread, charge)

Calculate charge density from Gaussian distributions.

Parameters:

Name Type Description Default
mu ndarray

Centers of Gaussian distributions

 required
spread ndarray

Spreads of Gaussian distributions

 required
charge ndarray

Charges of Gaussian distributions

 required

Returns:

Type Description
ndarray

Charge density on the grid
Source code in toolbox/calculator/elecpot.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
def calc_rho(
    self,
    mu: np.ndarray,
    spread: np.ndarray,
    charge: np.ndarray,
):
    """Calculate charge density from Gaussian distributions.

    Parameters
    ----------
    mu : np.ndarray
        Centers of Gaussian distributions
    spread : np.ndarray
        Spreads of Gaussian distributions
    charge : np.ndarray
        Charges of Gaussian distributions

    Returns
    -------
    np.ndarray
        Charge density on the grid
    """
    grid_edges = np.reshape(self.grid_edges, (1, -1))
    spread = np.reshape(spread, [-1, 1])
    mu = np.reshape(mu, [-1, 1])
    charge = np.reshape(charge, [-1, 1])

    # nat * ngrid
    out = gaussian_int(grid_edges[:, 1:], mu, spread) - gaussian_int(
        grid_edges[:, :-1], mu, spread
    )
    out = np.sum(out * charge, axis=0)
    # deal with periodic boundary
    # left image
    l_out = gaussian_int(grid_edges[:, 1:] - self.l_box, mu, spread) - gaussian_int(
        grid_edges[:, :-1] - self.l_box, mu, spread
    )
    l_out = np.sum(l_out * charge, axis=0)
    out += l_out
    # right image
    r_out = gaussian_int(grid_edges[:, 1:] + self.l_box, mu, spread) - gaussian_int(
        grid_edges[:, :-1] + self.l_box, mu, spread
    )
    r_out = np.sum(r_out * charge, axis=0)
    out += r_out

    np.testing.assert_almost_equal(out.sum(), charge.sum())

    rho = out / self.grid_volume
    return rho
run(mu=None, spread=None, charge=None, atoms=None)

Calculate electrostatic potential from Gaussian charge distributions.

Parameters:

Name Type Description Default
mu Optional[ndarray]

Centers of Gaussian distributions, by default None

 None
spread Optional[ndarray]

Spreads of Gaussian distributions, by default None

 None
charge Optional[ndarray]

Charges of Gaussian distributions, by default None

 None
atoms Optional[Atoms]

ASE Atoms object, by default None

 None

Returns:

Type Description
ndarray

Electrostatic potential
Source code in toolbox/calculator/elecpot.py 
289
290
291
292
293
294
295
296
297
298
299
300
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
def run(
    self,
    mu: Optional[np.ndarray] = None,
    spread: Optional[np.ndarray] = None,
    charge: Optional[np.ndarray] = None,
    atoms: Optional[Atoms] = None,
):
    """Calculate electrostatic potential from Gaussian charge distributions.

    Parameters
    ----------
    mu : Optional[np.ndarray], optional
        Centers of Gaussian distributions, by default None
    spread : Optional[np.ndarray], optional
        Spreads of Gaussian distributions, by default None
    charge : Optional[np.ndarray], optional
        Charges of Gaussian distributions, by default None
    atoms : Optional[Atoms], optional
        ASE Atoms object, by default None

    Returns
    -------
    np.ndarray
        Electrostatic potential
    """
    if mu is None:
        mu = atoms.get_positions()[:, 2]
    if spread is None:
        assert self.spread_dict is not None, "spread_dict is not set"
        spread = np.array([self.spread_dict[s] for s in atoms.symbols])
    if charge is None:
        if self.charge_dict is not None:
            charge = np.array([self.charge_dict[s] for s in atoms.symbols])
        else:
            charge = atoms.get_initial_charges()

    rho = self.calc_rho(mu, spread, charge)
    calculator = ElecPotentialCalculator(rho, self.grids)
    phi = calculator.calculate(l_box=self.l_box)
    return phi

WannierHartreePotentialCalculator

Bases: GaussianElecPotentialCalculator

Calculator for Hartree potential from Wannier functions.

This calculator extends GaussianElecPotentialCalculator to compute Hartree potential from Wannier functions.

Source code in toolbox/calculator/elecpot.py 
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
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
class WannierHartreePotentialCalculator(GaussianElecPotentialCalculator):
    """Calculator for Hartree potential from Wannier functions.

    This calculator extends GaussianElecPotentialCalculator to
    compute Hartree potential from Wannier functions.
    """

    def __init__(
        self,
        grids: np.ndarray,
        l_box: float,
        cross_area: float,
        spread_dict: Optional[dict[str, float]] = None,
        charge_dict: Optional[dict[str, float]] = None,
    ) -> None:
        """Initialize WannierHartreePotentialCalculator.

        Parameters
        ----------
        grids : np.ndarray
            Grid points in Ã…
        l_box : float
            Box length in Ã…
        cross_area : float
            Cross-sectional area in Ų
        spread_dict : Optional[Dict[str, float]], optional
            Dictionary of Gaussian spreads for each element, by default None
        charge_dict : Optional[Dict[str, float]], optional
            Dictionary of charges for each element, by default None
        """
        super().__init__(grids, l_box, cross_area, spread_dict, charge_dict)

    def run(
        self,
        mu: Optional[np.ndarray] = None,
        spread: Optional[np.ndarray] = None,
        charge: Optional[np.ndarray] = None,
        atoms: Optional[Atoms] = None,
        dname: Optional[str] = None,
        fname_coord: Optional[str] = "coord.xyz",
        fname_wannier: Optional[str] = "wannier.xyz",
    ):
        """Calculate Hartree potential from Wannier functions.

        Parameters
        ----------
        mu : Optional[np.ndarray], optional
            Centers of Gaussian distributions, by default None
        spread : Optional[np.ndarray], optional
            Spreads of Gaussian distributions, by default None
        charge : Optional[np.ndarray], optional
            Charges of Gaussian distributions, by default None
        atoms : Optional[Atoms], optional
            ASE Atoms object, by default None
        dname : Optional[str], optional
            Directory name containing coordinate files, by default None
        fname_coord : Optional[str], optional
            Coordinate file name, by default "coord.xyz"
        fname_wannier : Optional[str], optional
            Wannier coordinate file name, by default "wannier.xyz"

        Returns
        -------
        np.ndarray
            Negative of the electrostatic potential (Hartree potential)
        """
        if dname is not None:
            _atoms = io.read(os.path.join(dname, fname_coord))
            wannier_atoms = io.read(os.path.join(dname, fname_wannier))
            atoms = _atoms + wannier_atoms
            atoms.set_pbc(True)
            atoms.wrap()

        phi = super().run(mu, spread, charge, atoms)
        return -phi
__init__(grids, l_box, cross_area, spread_dict=None, charge_dict=None)

Initialize WannierHartreePotentialCalculator.

Parameters:

Name Type Description Default
grids ndarray

Grid points in Ã…

 required
l_box float

Box length in Ã…

 required
cross_area float

Cross-sectional area in Ų

 required
spread_dict Optional[Dict[str, float]]

Dictionary of Gaussian spreads for each element, by default None

 None
charge_dict Optional[Dict[str, float]]

Dictionary of charges for each element, by default None

 None
Source code in toolbox/calculator/elecpot.py 
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def __init__(
    self,
    grids: np.ndarray,
    l_box: float,
    cross_area: float,
    spread_dict: Optional[dict[str, float]] = None,
    charge_dict: Optional[dict[str, float]] = None,
) -> None:
    """Initialize WannierHartreePotentialCalculator.

    Parameters
    ----------
    grids : np.ndarray
        Grid points in Ã…
    l_box : float
        Box length in Ã…
    cross_area : float
        Cross-sectional area in Ų
    spread_dict : Optional[Dict[str, float]], optional
        Dictionary of Gaussian spreads for each element, by default None
    charge_dict : Optional[Dict[str, float]], optional
        Dictionary of charges for each element, by default None
    """
    super().__init__(grids, l_box, cross_area, spread_dict, charge_dict)
run(mu=None, spread=None, charge=None, atoms=None, dname=None, fname_coord='coord.xyz', fname_wannier='wannier.xyz')

Calculate Hartree potential from Wannier functions.

Parameters:

Name Type Description Default
mu Optional[ndarray]

Centers of Gaussian distributions, by default None

 None
spread Optional[ndarray]

Spreads of Gaussian distributions, by default None

 None
charge Optional[ndarray]

Charges of Gaussian distributions, by default None

 None
atoms Optional[Atoms]

ASE Atoms object, by default None

 None
dname Optional[str]

Directory name containing coordinate files, by default None

 None
fname_coord Optional[str]

Coordinate file name, by default "coord.xyz"

 'coord.xyz'
fname_wannier Optional[str]

Wannier coordinate file name, by default "wannier.xyz"

 'wannier.xyz'

Returns:

Type Description
ndarray

Negative of the electrostatic potential (Hartree potential)
Source code in toolbox/calculator/elecpot.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
def run(
    self,
    mu: Optional[np.ndarray] = None,
    spread: Optional[np.ndarray] = None,
    charge: Optional[np.ndarray] = None,
    atoms: Optional[Atoms] = None,
    dname: Optional[str] = None,
    fname_coord: Optional[str] = "coord.xyz",
    fname_wannier: Optional[str] = "wannier.xyz",
):
    """Calculate Hartree potential from Wannier functions.

    Parameters
    ----------
    mu : Optional[np.ndarray], optional
        Centers of Gaussian distributions, by default None
    spread : Optional[np.ndarray], optional
        Spreads of Gaussian distributions, by default None
    charge : Optional[np.ndarray], optional
        Charges of Gaussian distributions, by default None
    atoms : Optional[Atoms], optional
        ASE Atoms object, by default None
    dname : Optional[str], optional
        Directory name containing coordinate files, by default None
    fname_coord : Optional[str], optional
        Coordinate file name, by default "coord.xyz"
    fname_wannier : Optional[str], optional
        Wannier coordinate file name, by default "wannier.xyz"

    Returns
    -------
    np.ndarray
        Negative of the electrostatic potential (Hartree potential)
    """
    if dname is not None:
        _atoms = io.read(os.path.join(dname, fname_coord))
        wannier_atoms = io.read(os.path.join(dname, fname_wannier))
        atoms = _atoms + wannier_atoms
        atoms.set_pbc(True)
        atoms.wrap()

    phi = super().run(mu, spread, charge, atoms)
    return -phi

LAMMPS Calculator

toolbox.calculator.lammps

LAMMPS calculator module.

This module provides a calculator for running LAMMPS molecular dynamics simulations, extending the BashCalculator class.

LammpsCalculator

Bases: BashCalculator

A calculator for running LAMMPS molecular dynamics simulations.

This class extends BashCalculator to provide specific functionality for LAMMPS calculations, including completion checking.

Source code in toolbox/calculator/lammps.py 
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
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
class LammpsCalculator(BashCalculator):
    """A calculator for running LAMMPS molecular dynamics simulations.

    This class extends BashCalculator to provide specific functionality for
    LAMMPS calculations, including completion checking.
    """

    def __init__(self, work_dir) -> None:
        """Initialize LammpsCalculator.

        Parameters
        ----------
        work_dir : str
            The working directory where LAMMPS calculations will be executed
        """
        super().__init__(work_dir)

    def run(
        self,
        command: str = "lmp",
        stdin: str = "input.lmp",
        stdout: str = "lammps.stdout",
        stderr: str = "lammps.stderr",
        mpi_command: str = "mpiexec.hydra",
        ignore_finished_tag=False,
        modifier: str = None,
    ):
        """Run a LAMMPS calculation.

        Parameters
        ----------
        command : str, optional
            LAMMPS executable command, by default "lmp"
        stdin : str, optional
            Input file name, by default "input.lmp"
        stdout : str, optional
            Standard output file name, by default "lammps.stdout"
        stderr : str, optional
            Standard error file name, by default "lammps.stderr"
        mpi_command : str, optional
            MPI command to use, by default "mpiexec.hydra"
        ignore_finished_tag : bool, optional
            Whether to ignore existing finished tag, by default False
        modifier : str, optional
            Additional command modifiers, by default None

        Notes
        -----
        For more information on LAMMPS run options, see:
        https://docs.lammps.org/Run_options.html
        """
        stdin = "-i " + stdin
        if modifier is not None:
            stdin += f" {modifier} "
        super().run(command, stdin, stdout, stderr, mpi_command, ignore_finished_tag)

    @staticmethod
    def _make_finished_tag(stdout):
        """Create a finished tag after successful LAMMPS calculation.

        Parameters
        ----------
        stdout : str
            Name of LAMMPS output file to check for completion
        """
        with open(stdout, "rb") as f:
            offset = -50
            while True:
                f.seek(offset, 2)
                lines = f.readlines()
                if len(lines) >= 2:
                    last_line = lines[-1]
                    break
                offset *= 2

        pattern = re.compile(r"Total wall time")
        if pattern.search(last_line.decode()) is not None:
            with open(os.path.join("finished_tag"), "w") as f:
                pass
        else:
            warning_msg = "LAMMPS calculation does not finish!"
            logging.warning(warning_msg)
__init__(work_dir)

Initialize LammpsCalculator.

Parameters:

Name Type Description Default
work_dir str

The working directory where LAMMPS calculations will be executed

 required
Source code in toolbox/calculator/lammps.py 
22
23
24
25
26
27
28
29
30
def __init__(self, work_dir) -> None:
    """Initialize LammpsCalculator.

    Parameters
    ----------
    work_dir : str
        The working directory where LAMMPS calculations will be executed
    """
    super().__init__(work_dir)
run(command='lmp', stdin='input.lmp', stdout='lammps.stdout', stderr='lammps.stderr', mpi_command='mpiexec.hydra', ignore_finished_tag=False, modifier=None)

Run a LAMMPS calculation.

Parameters:

Name Type Description Default
command str

LAMMPS executable command, by default "lmp"

 'lmp'
stdin str

Input file name, by default "input.lmp"

 'input.lmp'
stdout str

Standard output file name, by default "lammps.stdout"

 'lammps.stdout'
stderr str

Standard error file name, by default "lammps.stderr"

 'lammps.stderr'
mpi_command str

MPI command to use, by default "mpiexec.hydra"

 'mpiexec.hydra'
ignore_finished_tag bool

Whether to ignore existing finished tag, by default False

 False
modifier str

Additional command modifiers, by default None

 None
Notes

For more information on LAMMPS run options, see: https://docs.lammps.org/Run_options.html

Source code in toolbox/calculator/lammps.py 
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 run(
    self,
    command: str = "lmp",
    stdin: str = "input.lmp",
    stdout: str = "lammps.stdout",
    stderr: str = "lammps.stderr",
    mpi_command: str = "mpiexec.hydra",
    ignore_finished_tag=False,
    modifier: str = None,
):
    """Run a LAMMPS calculation.

    Parameters
    ----------
    command : str, optional
        LAMMPS executable command, by default "lmp"
    stdin : str, optional
        Input file name, by default "input.lmp"
    stdout : str, optional
        Standard output file name, by default "lammps.stdout"
    stderr : str, optional
        Standard error file name, by default "lammps.stderr"
    mpi_command : str, optional
        MPI command to use, by default "mpiexec.hydra"
    ignore_finished_tag : bool, optional
        Whether to ignore existing finished tag, by default False
    modifier : str, optional
        Additional command modifiers, by default None

    Notes
    -----
    For more information on LAMMPS run options, see:
    https://docs.lammps.org/Run_options.html
    """
    stdin = "-i " + stdin
    if modifier is not None:
        stdin += f" {modifier} "
    super().run(command, stdin, stdout, stderr, mpi_command, ignore_finished_tag)

I/O Module

ASE Logger

toolbox.io.ase_logger

ASE MD Logger module.

This module provides an extended MDLogger class that can write trajectory files during molecular dynamics simulations.

MDLogger

Bases: MDLogger

Extended MD Logger for ASE molecular dynamics.

This class extends ASE's MDLogger to provide additional functionality for writing trajectory files during MD simulations.

Source code in toolbox/io/ase_logger.py 
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
class MDLogger(_MDLogger):
    """Extended MD Logger for ASE molecular dynamics.

    This class extends ASE's MDLogger to provide additional functionality
    for writing trajectory files during MD simulations.
    """

    def __init__(
        self,
        dyn: Any,  # not fully annotated so far to avoid a circular import
        atoms: Atoms,
        logfile: Union[IO, str],
        fname: str,
        write_kwargs: Optional[Dict] = None,
        **kwargs,
    ):
        if write_kwargs is None:
            write_kwargs = {}
        super().__init__(dyn, atoms, logfile, **kwargs)
        self.fname = fname
        self.write_kwargs = write_kwargs

    def __call__(self):
        """Write trajectory data to file.

        This method is called during MD simulation to write the current
        atomic configuration to the specified file.
        """
        super().__call__()
        self.atoms.write(self.fname, append=True, **self.write_kwargs)
__call__()

Write trajectory data to file.

This method is called during MD simulation to write the current atomic configuration to the specified file.

Source code in toolbox/io/ase_logger.py 
36
37
38
39
40
41
42
43
def __call__(self):
    """Write trajectory data to file.

    This method is called during MD simulation to write the current
    atomic configuration to the specified file.
    """
    super().__call__()
    self.atoms.write(self.fname, append=True, **self.write_kwargs)

BPNN I/O

toolbox.io.bpnn

BPNN data reader module.

This module provides functionality to read BPNN (Behler-Parrinello Neural Network) data files for machine learning potentials.

read_data(fname='input.data')

Read BPNN data file.

Parameters:

Name Type Description Default
fname str

Path to BPNN data file, by default "input.data"

 'input.data'

Returns:

Type Description
tuple

Tuple of (box, coord, charge, symbol, energy, force) where: - box: array of box parameters with shape (n_frames, 9) - coord: array of coordinates with shape (n_frames, n_atoms*3) - charge: array of charges with shape (n_frames, n_atoms) - symbol: list of element symbols for each frame - energy: array of energies with shape (n_frames,) - force: array of forces with shape (n_frames, n_atoms*3)
Source code in toolbox/io/bpnn.py 
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
66
67
68
69
70
71
72
73
74
75
76
77
def read_data(fname="input.data"):
    """Read BPNN data file.

    Parameters
    ----------
    fname : str, optional
        Path to BPNN data file, by default "input.data"

    Returns
    -------
    tuple
        Tuple of (box, coord, charge, symbol, energy, force) where:
        - box: array of box parameters with shape (n_frames, 9)
        - coord: array of coordinates with shape (n_frames, n_atoms*3)
        - charge: array of charges with shape (n_frames, n_atoms)
        - symbol: list of element symbols for each frame
        - energy: array of energies with shape (n_frames,)
        - force: array of forces with shape (n_frames, n_atoms*3)
    """
    box = []
    coord = []
    charge = []
    symbol = []
    energy = []
    force = []

    flag = False
    count = 0
    _symbol = []
    with open(fname, encoding="UTF-8") as f:
        lines = f.readlines()
        for line in lines:
            line = line.strip()
            if line == "begin":
                flag = True
                count = count + 1
            if line == "end":
                flag = False
                symbol.append(_symbol)
                _symbol = []
            if flag is False:
                continue
            line = line.split()
            if line[0] == "lattice":
                box.append(line[1:])
            if line[0] == "atom":
                coord.append(line[1:4])
                charge.append(line[5])
                force.append(line[7:10])
                _symbol.append(line[4])
            if line[0] == "energy":
                energy.append(line[1])

    box = np.array(box, dtype=np.float64)
    if len(box) > 0:
        box = np.reshape(box, (count, 9)) * AU_TO_ANG
    charge = np.array(charge, dtype=np.float64)
    charge = np.reshape(charge, (count, -1))
    coord = np.array(coord, dtype=np.float64)
    coord = np.reshape(coord, (count, -1)) * AU_TO_ANG
    energy = np.array(energy, dtype=np.float64) * AU_TO_EV
    force = np.array(force, dtype=np.float64)
    force = np.reshape(force, (count, -1)) * AU_TO_EV / AU_TO_ANG

    return box, coord, charge, symbol, energy, force

System Builder

toolbox.io.build

Molecular system building module.

This module provides classes and functions for building molecular systems, including solution boxes, interfaces, and adding ions or molecules to existing systems.

ElectrolyteBox

Bases: WaterBox

Class for electrolyte solution box configuration.

This class extends WaterBox to handle uniform electrolyte solutions with specified solutes and concentrations.

Parameters:

Name Type Description Default
boundary Union[List, ndarray]

Box dimensions as [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]

 required
solutes Dict[str, float]

Dictionary mapping solute names to concentrations in mol/L

 required
Source code in toolbox/io/build.py 
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
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
class ElectrolyteBox(WaterBox):
    """Class for electrolyte solution box configuration.

    This class extends WaterBox to handle uniform electrolyte
    solutions with specified solutes and concentrations.

    Parameters
    ----------
    boundary : Union[List, np.ndarray]
        Box dimensions as [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]
    solutes : Dict[str, float]
        Dictionary mapping solute names to concentrations in mol/L
    """

    def __init__(
        self,
        boundary: Union[List, np.ndarray],
        solutes: dict[str, float],
        **kwargs,
    ) -> None:
        """Initialize ElectrolyteBox.

        Parameters
        ----------
        boundary : Union[List, np.ndarray]
            Box dimensions
        solutes : Dict[str, float]
            Dictionary mapping solute names to concentrations in mol/L
        **kwargs
            Additional keyword arguments passed to parent class
        """
        super().__init__(boundary, **kwargs)
        self.solutes = solutes

    def write(
        self,
        fname,
        n_wat: Optional[int] = None,
        seed: int = -1,
        verbose=False,
        **kwargs,
    ):
        """Write electrolyte box configuration to file.

        Parameters
        ----------
        fname : str
            Output filename
        n_wat : Optional[int], optional
            Number of water molecules. If None, uses calculated value from density
        seed : int, optional
            Random seed for reproducible molecule placement. Use -1 for random behavior
        verbose : bool, optional
            If True, keeps temporary files, by default False
        **kwargs
            Additional keyword arguments passed to ase.io.write
        """
        atoms = build.molecule("H2O")
        io.write("tmp.pdb", atoms)
        water = mda.Universe("tmp.pdb")

        if n_wat is None:
            n_wat = self.n_wat

        structures = [
            mdapackmol.PackmolStructure(
                water,
                number=n_wat,
                instructions=[f"inside box {self.boundary_string}", f"seed {seed}"],
            )
        ]
        for k, v in self.solutes.items():
            fnames = glob.glob(os.path.join(os.path.dirname(__file__), f"ion_structure_lib/*/{k}.*"))
            if len(fnames) > 0:
                atoms = io.read(fnames[0])
            else:
                atoms = Atoms(k, positions=[[0, 0, 0]])
            io.write("tmp.pdb", atoms)
            solute = mda.Universe("tmp.pdb")
            structures.append(
                mdapackmol.PackmolStructure(
                    solute,
                    number=int(n_wat / 55.6 * v),
                    instructions=[f"inside box {self.boundary_string}", f"seed {seed}"],
                )
            )

        system = mdapackmol.packmol(structures)
        system.atoms.write(fname, **kwargs)
        if not verbose:
            os.remove("tmp.pdb")
            os.remove("packmol.stdout")
__init__(boundary, solutes, **kwargs)

Initialize ElectrolyteBox.

Parameters:

Name Type Description Default
boundary Union[List, ndarray]

Box dimensions

 required
solutes Dict[str, float]

Dictionary mapping solute names to concentrations in mol/L

 required
**kwargs

Additional keyword arguments passed to parent class

 {}
Source code in toolbox/io/build.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
def __init__(
    self,
    boundary: Union[List, np.ndarray],
    solutes: dict[str, float],
    **kwargs,
) -> None:
    """Initialize ElectrolyteBox.

    Parameters
    ----------
    boundary : Union[List, np.ndarray]
        Box dimensions
    solutes : Dict[str, float]
        Dictionary mapping solute names to concentrations in mol/L
    **kwargs
        Additional keyword arguments passed to parent class
    """
    super().__init__(boundary, **kwargs)
    self.solutes = solutes
write(fname, n_wat=None, seed=-1, verbose=False, **kwargs)

Write electrolyte box configuration to file.

Parameters:

Name Type Description Default
fname str

Output filename

 required
n_wat Optional[int]

Number of water molecules. If None, uses calculated value from density

 None
seed int

Random seed for reproducible molecule placement. Use -1 for random behavior

 -1
verbose bool

If True, keeps temporary files, by default False

 False
**kwargs

Additional keyword arguments passed to ase.io.write

 {}
Source code in toolbox/io/build.py 
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
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
def write(
    self,
    fname,
    n_wat: Optional[int] = None,
    seed: int = -1,
    verbose=False,
    **kwargs,
):
    """Write electrolyte box configuration to file.

    Parameters
    ----------
    fname : str
        Output filename
    n_wat : Optional[int], optional
        Number of water molecules. If None, uses calculated value from density
    seed : int, optional
        Random seed for reproducible molecule placement. Use -1 for random behavior
    verbose : bool, optional
        If True, keeps temporary files, by default False
    **kwargs
        Additional keyword arguments passed to ase.io.write
    """
    atoms = build.molecule("H2O")
    io.write("tmp.pdb", atoms)
    water = mda.Universe("tmp.pdb")

    if n_wat is None:
        n_wat = self.n_wat

    structures = [
        mdapackmol.PackmolStructure(
            water,
            number=n_wat,
            instructions=[f"inside box {self.boundary_string}", f"seed {seed}"],
        )
    ]
    for k, v in self.solutes.items():
        fnames = glob.glob(os.path.join(os.path.dirname(__file__), f"ion_structure_lib/*/{k}.*"))
        if len(fnames) > 0:
            atoms = io.read(fnames[0])
        else:
            atoms = Atoms(k, positions=[[0, 0, 0]])
        io.write("tmp.pdb", atoms)
        solute = mda.Universe("tmp.pdb")
        structures.append(
            mdapackmol.PackmolStructure(
                solute,
                number=int(n_wat / 55.6 * v),
                instructions=[f"inside box {self.boundary_string}", f"seed {seed}"],
            )
        )

    system = mdapackmol.packmol(structures)
    system.atoms.write(fname, **kwargs)
    if not verbose:
        os.remove("tmp.pdb")
        os.remove("packmol.stdout")

Interface

Class for creating interface systems.

This class creates interface systems between a slab and water box for electrochemical simulations.

Source code in toolbox/io/build.py 
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
300
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
class Interface:
    """Class for creating interface systems.

    This class creates interface systems between a slab
    and water box for electrochemical simulations.
    """

    def __init__(
        self,
        slab: Atoms,
        l_water: float = 30,
    ) -> None:
        """Initialize Interface.

        Parameters
        ----------
        slab : ase.Atoms
            Slab atoms object
        l_water : float, optional
            Length of water box in Ã…, by default 30
        """
        # shift half cell and wrap
        coord = slab.get_positions()
        z = coord[:, 2]
        l_slab = z.max() - z.min()
        a = slab.get_cell()[0][0]
        b = slab.get_cell()[1][1]
        c = l_slab + l_water
        new_cell = [a, b, c]
        shift_z = -z.min() - l_slab / 2
        coord[:, 2] += shift_z
        slab = Atoms(
            slab.get_chemical_symbols(), positions=coord, cell=new_cell, pbc=True
        )
        slab.wrap()

        self.slab = slab
        self.boundary = [[0, a], [0, b], [l_slab / 2, l_slab / 2 + l_water]]

    def run(
        self,
        rho: float = 1.0,
        n_wat: Optional[int] = None,
        seed: int = -1,
        sol: Optional[SolutionBox] = None,
        verbose=False,
    ) -> Atoms:
        """Create interface system.

        Parameters
        ----------
        rho : float, optional
            Water density in g/cm³, by default 1.0
        n_wat : Optional[int], optional
            Number of water molecules, by default None
        seed : int, optional
            Random seed, by default -1
        sol : Optional[SolutionBox], optional
            Solution box object, by default None
        verbose : bool, optional
            If True, keeps temporary files, by default False

        Returns
        -------
        ase.Atoms
            Combined interface system
        """
        if sol is None:
            sol = WaterBox(
                rho=rho,
                boundary=self.boundary,
                slit=[1.0, 1.0, 2.5],
            )
        sol.write("waterbox.xyz", n_wat=n_wat, verbose=verbose, seed=seed)
        waterbox = io.read("waterbox.xyz")
        waterbox.set_cell(self.slab.get_cell())
        waterbox.set_pbc(True)
        waterbox.center(axis=2)

        self.atoms = waterbox + self.slab
        # self.atoms.set_cell(self.slab.get_cell())
        self.atoms.set_pbc(True)
        if not verbose:
            os.remove("waterbox.xyz")
        return self.atoms
__init__(slab, l_water=30)

Initialize Interface.

Parameters:

Name Type Description Default
slab Atoms

Slab atoms object

 required
l_water float

Length of water box in Ã…, by default 30

 30
Source code in toolbox/io/build.py 
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
def __init__(
    self,
    slab: Atoms,
    l_water: float = 30,
) -> None:
    """Initialize Interface.

    Parameters
    ----------
    slab : ase.Atoms
        Slab atoms object
    l_water : float, optional
        Length of water box in Ã…, by default 30
    """
    # shift half cell and wrap
    coord = slab.get_positions()
    z = coord[:, 2]
    l_slab = z.max() - z.min()
    a = slab.get_cell()[0][0]
    b = slab.get_cell()[1][1]
    c = l_slab + l_water
    new_cell = [a, b, c]
    shift_z = -z.min() - l_slab / 2
    coord[:, 2] += shift_z
    slab = Atoms(
        slab.get_chemical_symbols(), positions=coord, cell=new_cell, pbc=True
    )
    slab.wrap()

    self.slab = slab
    self.boundary = [[0, a], [0, b], [l_slab / 2, l_slab / 2 + l_water]]
run(rho=1.0, n_wat=None, seed=-1, sol=None, verbose=False)

Create interface system.

Parameters:

Name Type Description Default
rho float

Water density in g/cm³, by default 1.0

 1.0
n_wat Optional[int]

Number of water molecules, by default None

 None
seed int

Random seed, by default -1

 -1
sol Optional[SolutionBox]

Solution box object, by default None

 None
verbose bool

If True, keeps temporary files, by default False

 False

Returns:

Type Description
Atoms

Combined interface system
Source code in toolbox/io/build.py 
289
290
291
292
293
294
295
296
297
298
299
300
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
def run(
    self,
    rho: float = 1.0,
    n_wat: Optional[int] = None,
    seed: int = -1,
    sol: Optional[SolutionBox] = None,
    verbose=False,
) -> Atoms:
    """Create interface system.

    Parameters
    ----------
    rho : float, optional
        Water density in g/cm³, by default 1.0
    n_wat : Optional[int], optional
        Number of water molecules, by default None
    seed : int, optional
        Random seed, by default -1
    sol : Optional[SolutionBox], optional
        Solution box object, by default None
    verbose : bool, optional
        If True, keeps temporary files, by default False

    Returns
    -------
    ase.Atoms
        Combined interface system
    """
    if sol is None:
        sol = WaterBox(
            rho=rho,
            boundary=self.boundary,
            slit=[1.0, 1.0, 2.5],
        )
    sol.write("waterbox.xyz", n_wat=n_wat, verbose=verbose, seed=seed)
    waterbox = io.read("waterbox.xyz")
    waterbox.set_cell(self.slab.get_cell())
    waterbox.set_pbc(True)
    waterbox.center(axis=2)

    self.atoms = waterbox + self.slab
    # self.atoms.set_cell(self.slab.get_cell())
    self.atoms.set_pbc(True)
    if not verbose:
        os.remove("waterbox.xyz")
    return self.atoms

SolutionBox

Base class for solution box configuration.

This class defines the boundary conditions for solution systems with optional slit geometry.

Parameters:

Name Type Description Default
boundary Union[List, ndarray]

Box dimensions as [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]

 required
slit Union[float, List, ndarray]

Slit distance or dimensions, by default 1.0

 1.0
Source code in toolbox/io/build.py 
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
66
67
68
69
70
class SolutionBox:
    """Base class for solution box configuration.

    This class defines the boundary conditions for solution
    systems with optional slit geometry.

    Parameters
    ----------
    boundary : Union[List, np.ndarray]
        Box dimensions as [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]
    slit : Union[float, List, np.ndarray], optional
        Slit distance or dimensions, by default 1.0
    """

    def __init__(
        self,
        boundary: Union[List, np.ndarray],
        slit: Union[float, List, np.ndarray] = 1.0,
    ) -> None:
        """Initialize SolutionBox.

        Parameters
        ----------
        boundary : Union[List, np.ndarray]
            Box dimensions
        slit : Union[float, List, np.ndarray], optional
            Slit distance or dimensions, by default 1.0
        """
        boundary = np.array(boundary)
        assert len(boundary) == 3
        _boundary = np.reshape(boundary, (3, -1))
        # [[a1, a2], [b1, b2], [c1, c2]]
        if len(_boundary[0]) == 1:
            self.boundary = np.concatenate([np.zeros((3, 1)), _boundary], axis=-1)
        elif len(_boundary[0]) == 2:
            self.boundary = _boundary
        else:
            raise AttributeError(
                "boundary must be [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]"
            )

        slit = np.reshape(slit, (-1))
        b = self.boundary.copy()
        if (len(slit) == 1) or (len(slit) == 3):
            b[:, 0] += slit
            b[:, 1] -= slit
        else:
            raise AttributeError("")
        self.boundary_string = np.array2string(np.transpose(b).flatten())[1:-1]
__init__(boundary, slit=1.0)

Initialize SolutionBox.

Parameters:

Name Type Description Default
boundary Union[List, ndarray]

Box dimensions

 required
slit Union[float, List, ndarray]

Slit distance or dimensions, by default 1.0

 1.0
Source code in toolbox/io/build.py 
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
def __init__(
    self,
    boundary: Union[List, np.ndarray],
    slit: Union[float, List, np.ndarray] = 1.0,
) -> None:
    """Initialize SolutionBox.

    Parameters
    ----------
    boundary : Union[List, np.ndarray]
        Box dimensions
    slit : Union[float, List, np.ndarray], optional
        Slit distance or dimensions, by default 1.0
    """
    boundary = np.array(boundary)
    assert len(boundary) == 3
    _boundary = np.reshape(boundary, (3, -1))
    # [[a1, a2], [b1, b2], [c1, c2]]
    if len(_boundary[0]) == 1:
        self.boundary = np.concatenate([np.zeros((3, 1)), _boundary], axis=-1)
    elif len(_boundary[0]) == 2:
        self.boundary = _boundary
    else:
        raise AttributeError(
            "boundary must be [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]"
        )

    slit = np.reshape(slit, (-1))
    b = self.boundary.copy()
    if (len(slit) == 1) or (len(slit) == 3):
        b[:, 0] += slit
        b[:, 1] -= slit
    else:
        raise AttributeError("")
    self.boundary_string = np.array2string(np.transpose(b).flatten())[1:-1]

WaterBox

Bases: SolutionBox

Class for water box configuration.

This class extends SolutionBox to specifically handle water molecules with density calculations.

Parameters:

Name Type Description Default
boundary Union[List, ndarray]

Box dimensions as [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]

 required
slit Union[float, List, ndarray]

Slit distance or dimensions, by default 1.0

 1.0
rho float

Water density in g/cm³

 1.0
Source code in toolbox/io/build.py 
 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
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
class WaterBox(SolutionBox):
    """Class for water box configuration.

    This class extends SolutionBox to specifically handle
    water molecules with density calculations.

    Parameters
    ----------
    boundary : Union[List, np.ndarray]
        Box dimensions as [a, b, c] or [[a1, a2], [b1, b2], [c1, c2]]
    slit : Union[float, List, np.ndarray], optional
        Slit distance or dimensions, by default 1.0
    rho : float
        Water density in g/cm³
    """

    def __init__(
        self,
        boundary: Union[List, np.ndarray],
        slit: Union[float, List, np.ndarray] = 1.0,
        rho: float = 1.0,
    ) -> None:
        """Initialize WaterBox.

        Parameters
        ----------
        boundary : Union[List, np.ndarray]
            Box dimensions
        slit : Union[float, List, np.ndarray], optional
            Slit distance or dimensions, by default 1.0
        rho : float, optional
            Water density in g/cm³, by default 1.0
        """
        super().__init__(boundary, slit)

        volume = np.prod(np.diff(self.boundary, axis=-1))
        self.n_wat = calc_water_number(rho, volume)

    def write(
        self,
        fname,
        n_wat: Optional[int] = None,
        seed: int = -1,
        verbose: bool = False,
        **kwargs,
    ):
        """Write water box configuration to file.

        Parameters
        ----------
        fname : str
            Output filename
        n_wat : Optional[int], optional
            Number of water molecules. If None, uses calculated value from density
        seed : int, optional
            Random seed for reproducible molecule placement. Use -1 for random behavior
        verbose : bool, optional
            If True, keeps temporary files, by default False
        **kwargs
            Additional keyword arguments passed to ase.io.write
        """
        atoms = build.molecule("H2O")
        io.write("water.pdb", atoms)
        water = mda.Universe("water.pdb")

        if n_wat is None:
            n_wat = self.n_wat

        system = mdapackmol.packmol(
            [
                mdapackmol.PackmolStructure(
                    water,
                    number=n_wat,
                    instructions=[f"inside box {self.boundary_string}", f"seed {seed}"],
                )
            ]
        )
        system.atoms.write(fname, **kwargs)
        if not verbose:
            os.remove("water.pdb")
            os.remove("packmol.stdout")
__init__(boundary, slit=1.0, rho=1.0)

Initialize WaterBox.

Parameters:

Name Type Description Default
boundary Union[List, ndarray]

Box dimensions

 required
slit Union[float, List, ndarray]

Slit distance or dimensions, by default 1.0

 1.0
rho float

Water density in g/cm³, by default 1.0

 1.0
Source code in toolbox/io/build.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def __init__(
    self,
    boundary: Union[List, np.ndarray],
    slit: Union[float, List, np.ndarray] = 1.0,
    rho: float = 1.0,
) -> None:
    """Initialize WaterBox.

    Parameters
    ----------
    boundary : Union[List, np.ndarray]
        Box dimensions
    slit : Union[float, List, np.ndarray], optional
        Slit distance or dimensions, by default 1.0
    rho : float, optional
        Water density in g/cm³, by default 1.0
    """
    super().__init__(boundary, slit)

    volume = np.prod(np.diff(self.boundary, axis=-1))
    self.n_wat = calc_water_number(rho, volume)
write(fname, n_wat=None, seed=-1, verbose=False, **kwargs)

Write water box configuration to file.

Parameters:

Name Type Description Default
fname str

Output filename

 required
n_wat Optional[int]

Number of water molecules. If None, uses calculated value from density

 None
seed int

Random seed for reproducible molecule placement. Use -1 for random behavior

 -1
verbose bool

If True, keeps temporary files, by default False

 False
**kwargs

Additional keyword arguments passed to ase.io.write

 {}
Source code in toolbox/io/build.py 
111
112
113
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
def write(
    self,
    fname,
    n_wat: Optional[int] = None,
    seed: int = -1,
    verbose: bool = False,
    **kwargs,
):
    """Write water box configuration to file.

    Parameters
    ----------
    fname : str
        Output filename
    n_wat : Optional[int], optional
        Number of water molecules. If None, uses calculated value from density
    seed : int, optional
        Random seed for reproducible molecule placement. Use -1 for random behavior
    verbose : bool, optional
        If True, keeps temporary files, by default False
    **kwargs
        Additional keyword arguments passed to ase.io.write
    """
    atoms = build.molecule("H2O")
    io.write("water.pdb", atoms)
    water = mda.Universe("water.pdb")

    if n_wat is None:
        n_wat = self.n_wat

    system = mdapackmol.packmol(
        [
            mdapackmol.PackmolStructure(
                water,
                number=n_wat,
                instructions=[f"inside box {self.boundary_string}", f"seed {seed}"],
            )
        ]
    )
    system.atoms.write(fname, **kwargs)
    if not verbose:
        os.remove("water.pdb")
        os.remove("packmol.stdout")

add_ion(atoms, ion, region, cutoff=2.0, max_trial=500)

Add an ion to atoms system in specified region.

This function attempts to place an ion in the specified region without overlapping with existing atoms.

Parameters:

Name Type Description Default
atoms Atoms

Existing atoms system

 required
ion Atoms

Ion to add (single atom)

 required
region list

Region boundaries as [x_min, x_max, y_min, y_max, z_min, z_max]

 required
cutoff float

Minimum distance from existing atoms, by default 2.0

 2.0
max_trial int

Maximum number of placement attempts, by default 500

 500

Returns:

Type Description
Atoms or None

New atoms system with ion added, or None if placement failed

Examples:

>>> from ase import io
>>> atoms = io.read("coord.xyz")
>>> ion = Atoms("K")
>>> atoms = add_ion(atoms, ion, [8.0, 13.0])
>>> ion = io.read("ClO4.pdb")
>>> atoms = add_ion(atoms, ion, [8.0, 13.0])
>>> print(atoms)
Source code in toolbox/io/build.py 
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
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
def add_ion(atoms, ion, region, cutoff=2.0, max_trial=500):
    """Add an ion to atoms system in specified region.

    This function attempts to place an ion in the specified region
    without overlapping with existing atoms.

    Parameters
    ----------
    atoms : ase.Atoms
        Existing atoms system
    ion : ase.Atoms
        Ion to add (single atom)
    region : list
        Region boundaries as [x_min, x_max, y_min, y_max, z_min, z_max]
    cutoff : float, optional
        Minimum distance from existing atoms, by default 2.0
    max_trial : int, optional
        Maximum number of placement attempts, by default 500

    Returns
    -------
    ase.Atoms or None
        New atoms system with ion added, or None if placement failed

    Examples
    --------
    >>> from ase import io
    >>> atoms = io.read("coord.xyz")
    >>> ion = Atoms("K")
    >>> atoms = add_ion(atoms, ion, [8.0, 13.0])
    >>> ion = io.read("ClO4.pdb")
    >>> atoms = add_ion(atoms, ion, [8.0, 13.0])
    >>> print(atoms)
    """
    coords = ion.get_positions()
    cog = np.mean(coords, axis=0)
    coords -= cog.reshape(1, 3)
    rotation_matrix = random_rotation_matrix()
    coords = np.dot(coords, rotation_matrix)

    flag = False
    for _ in range(max_trial):
        random_positions = get_region_random_location(atoms, region)
        random_positions = coords + random_positions.reshape(1, 3)
        ds = distance_array(
            random_positions, atoms.get_positions(), box=atoms.cell.cellpar()
        )
        if ds.min() > cutoff:
            flag = True
            break
    if flag:
        ion.set_positions(random_positions)
        new_atoms = atoms.copy()
        new_atoms.extend(ion)
        return new_atoms
    else:
        raise Warning("Failed to add ion")
        return None

add_water(atoms, region, cutoff=2.0, max_trial=500)

Add a water molecule to atoms system in specified region.

Parameters:

Name Type Description Default
atoms Atoms

Existing atoms system

 required
region list

Region boundaries as [x_min, x_max, y_min, y_max, z_min, z_max]

 required
cutoff float

Minimum distance from existing atoms, by default 2.0

 2.0
max_trial int

Maximum number of placement attempts, by default 500

 500

Returns:

Type Description
Atoms or None

New atoms system with water added, or None if placement failed
Source code in toolbox/io/build.py 
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def add_water(atoms, region, cutoff=2.0, max_trial=500):
    """Add a water molecule to atoms system in specified region.

    Parameters
    ----------
    atoms : ase.Atoms
        Existing atoms system
    region : list
        Region boundaries as [x_min, x_max, y_min, y_max, z_min, z_max]
    cutoff : float, optional
        Minimum distance from existing atoms, by default 2.0
    max_trial : int, optional
        Maximum number of placement attempts, by default 500

    Returns
    -------
    ase.Atoms or None
        New atoms system with water added, or None if placement failed
    """
    water = build.molecule("H2O")
    return add_ion(atoms, water, region, cutoff=cutoff, max_trial=max_trial)

get_region_random_location(atoms, region, extent=0.9)

Generate random location within specified region.

Parameters:

Name Type Description Default
atoms Atoms

Atoms system for box dimensions

 required
region list

Region boundaries as [x_min, x_max, y_min, y_max, z_min, z_max]

 required
extent float

Fraction of box extent to use, by default 0.9

 0.9

Returns:

Type Description
ndarray

Random position [x, y, z] within region
Source code in toolbox/io/build.py 
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
def get_region_random_location(atoms, region, extent=0.9):
    """Generate random location within specified region.

    Parameters
    ----------
    atoms : ase.Atoms
        Atoms system for box dimensions
    region : list
        Region boundaries as [x_min, x_max, y_min, y_max, z_min, z_max]
    extent : float, optional
        Fraction of box extent to use, by default 0.9

    Returns
    -------
    np.ndarray
        Random position [x, y, z] within region
    """
    x_region = [atoms.get_cell()[0][0] * (1 - extent), atoms.get_cell()[0][0] * extent]
    y_region = [atoms.get_cell()[1][1] * (1 - extent), atoms.get_cell()[1][1] * extent]

    location_x = np.random.uniform(x_region[0], x_region[1])
    location_y = np.random.uniform(y_region[0], y_region[1])
    location_z = np.random.uniform(region[0], region[1])

    return np.array([location_x, location_y, location_z])

is_atom_overlap(atoms, index, atom_num=1, r=1.5)

Check if atoms overlap with existing atoms.

This function checks if specified atoms in a molecule or ion overlap with existing atoms in the system.

Parameters:

Name Type Description Default
atoms Atoms

Existing atoms system

 required
index int

Index of first atom in the added molecule/ion

 required
atom_num int

Number of atoms in the added molecule/ion, by default 1

 1
r float

Minimum distance between atoms, recommended 1.0-1.8, by default 1.5

 1.5

Returns:

Type Description
bool

True if atoms overlap, False otherwise
Source code in toolbox/io/build.py 
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
def is_atom_overlap(atoms, index, atom_num=1, r=1.5):
    """Check if atoms overlap with existing atoms.

    This function checks if specified atoms in a molecule or ion
    overlap with existing atoms in the system.

    Parameters
    ----------
    atoms : ase.Atoms
        Existing atoms system
    index : int
        Index of first atom in the added molecule/ion
    atom_num : int, optional
        Number of atoms in the added molecule/ion, by default 1
    r : float, optional
        Minimum distance between atoms, recommended 1.0-1.8, by default 1.5

    Returns
    -------
    bool
        True if atoms overlap, False otherwise
    """
    distance = atoms.get_all_distances(mic=True)[index]
    min_index = np.argsort(distance)[atom_num]

    overlap = distance[min_index] < r
    return overlap

random_rotation_matrix()

Generate a random 3D rotation matrix.

Returns:

Type Description
ndarray

3x3 rotation matrix
Source code in toolbox/io/build.py 
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
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
def random_rotation_matrix():
    """Generate a random 3D rotation matrix.

    Returns
    -------
    np.ndarray
        3x3 rotation matrix
    """
    random_rotvec = np.random.randn(3)
    random_rotvec /= np.linalg.norm(random_rotvec)
    angle = np.random.rand() * 2 * np.pi
    cos_theta = np.cos(angle)
    sin_theta = np.sin(angle)
    one_minus_cos_theta = 1 - cos_theta

    rot_matrix = np.zeros((3, 3))
    rot_matrix[0, 0] = cos_theta + random_rotvec[0] ** 2 * one_minus_cos_theta
    rot_matrix[1, 1] = cos_theta + random_rotvec[1] ** 2 * one_minus_cos_theta
    rot_matrix[2, 2] = cos_theta + random_rotvec[2] ** 2 * one_minus_cos_theta

    rot_matrix[0, 1] = (
        random_rotvec[0] * random_rotvec[1] * one_minus_cos_theta
        - random_rotvec[2] * sin_theta
    )
    rot_matrix[1, 0] = (
        random_rotvec[0] * random_rotvec[1] * one_minus_cos_theta
        + random_rotvec[2] * sin_theta
    )

    rot_matrix[0, 2] = (
        random_rotvec[0] * random_rotvec[2] * one_minus_cos_theta
        + random_rotvec[1] * sin_theta
    )
    rot_matrix[2, 0] = (
        random_rotvec[0] * random_rotvec[2] * one_minus_cos_theta
        - random_rotvec[1] * sin_theta
    )

    rot_matrix[1, 2] = (
        random_rotvec[1] * random_rotvec[2] * one_minus_cos_theta
        - random_rotvec[0] * sin_theta
    )
    rot_matrix[2, 1] = (
        random_rotvec[1] * random_rotvec[2] * one_minus_cos_theta
        + random_rotvec[0] * sin_theta
    )

    return rot_matrix

CP2K I/O

toolbox.io.cp2k

CP2K input/output module.

This module provides classes for generating CP2K input files and parsing CP2K output files, including energies, forces, and various properties.

Cp2kCube

Class for handling CP2K cube files.

This class provides methods to read and analyze cube files generated by CP2K, including electron density, potential, and other volumetric data.

Parameters:

Name Type Description Default
fname str

Path to cube file

 required
Source code in toolbox/io/cp2k.py 
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
class Cp2kCube:
    """
    Class for handling CP2K cube files.

    This class provides methods to read and analyze cube files generated by CP2K,
    including electron density, potential, and other volumetric data.

    Parameters
    ----------
    fname : str
        Path to cube file
    """

    def __init__(self, fname) -> None:
        self.cube_data, self.atoms = read_cube_data(fname)
        self.n_grid = np.array(self.cube_data.shape)
        # 3 * 3
        self.cube_vectors = self.atoms.get_cell() / self.n_grid.reshape(3, 1)
        # cube volume [A^3]
        self.cube_volume = self.atoms.get_volume() / np.prod(self.n_grid)

        # cell_params = self.atoms.cell.cellpar()
        # try:
        #     assert not (False in (90. == cell_params[-3:]))
        #     self.cell_params = cell_params[:3]
        # except:
        #     raise ValueError("Cell is not orthogonal")

    # @property
    # def cube_grids(self):
    #     cube_grids = []
    #     for ii in range(3):
    #         cube_grids.append(
    #             np.arange(
    #                 0, self.cell_params[ii],
    #                 self.cell_params[ii] / self.cube_data.shape[ii]
    #             )[:self.cube_data.shape[ii]]
    #         )
    #     return cube_grids

    @property
    def mesh(self):
        """Mesh for cube data."""
        # generate mesh from self.cube_vectors
        x = np.arange(self.n_grid[0])
        y = np.arange(self.n_grid[1])
        z = np.arange(self.n_grid[2])
        coeff = np.meshgrid(x, y, z, indexing="ij")
        coeff_x = coeff[0].reshape(self.n_grid[0], self.n_grid[1], self.n_grid[2], 1)
        coeff_y = coeff[1].reshape(self.n_grid[0], self.n_grid[1], self.n_grid[2], 1)
        coeff_z = coeff[2].reshape(self.n_grid[0], self.n_grid[1], self.n_grid[2], 1)
        unit_x = self.cube_vectors[0].reshape(1, 1, 1, 3)
        unit_y = self.cube_vectors[1].reshape(1, 1, 1, 3)
        unit_z = self.cube_vectors[2].reshape(1, 1, 1, 3)
        # dimx * dimy * dimz * 3
        mesh = coeff_x * unit_x + coeff_y * unit_y + coeff_z * unit_z
        return mesh

    @property
    def dipole(self):
        """Dipole moment [e A]."""
        # cube_data: charge density [e / bohr^3]
        charge = self.cube_data * (self.cube_volume / AU_TO_ANG**3)
        charge = np.reshape(charge, [self.n_grid[0], self.n_grid[1], self.n_grid[2], 1])
        dipole = np.sum(self.mesh * charge, axis=(0, 1, 2))
        # in cp2k, the electron density is positive
        return -dipole

    def get_ave_cube(self, axis=2, gaussian_sigma=0.0):
        """Get averaged cube data along specified axis.

        Parameters
        ----------
        axis : int, optional
            Axis along which to average (0=x, 1=y, 2=z), by default 2
        gaussian_sigma : float, optional
            Standard deviation for Gaussian smoothing, by default 0.0

        Returns
        -------
        tuple
            Tuple containing (grid, averaged_data, smoothed_data)
        """
        if (
            hasattr(self, "axis")
            and self.axis == axis
            and hasattr(self, "ave_cube_data")
        ):
            pass
        else:
            self.axis = axis
            self.ave_grid = self.cube_vectors[self.axis][self.axis] * np.arange(
                self.n_grid[self.axis]
            )
            ave_axis = tuple(np.delete(np.arange(3), self.axis).tolist())
            self.ave_cube_data = np.mean(self.cube_data, axis=ave_axis)

        if gaussian_sigma > 0.0:
            self.ave_cube_data_convolve = gaussian_convolve(
                self.ave_grid, self.ave_cube_data, gaussian_sigma
            )
        else:
            self.ave_cube_data_convolve = copy.deepcopy(self.ave_cube_data)

        return (self.ave_grid, self.ave_cube_data, self.ave_cube_data_convolve)
dipole property

Dipole moment [e A].

mesh property

Mesh for cube data.

get_ave_cube(axis=2, gaussian_sigma=0.0)

Get averaged cube data along specified axis.

Parameters:

Name Type Description Default
axis int

Axis along which to average (0=x, 1=y, 2=z), by default 2

 2
gaussian_sigma float

Standard deviation for Gaussian smoothing, by default 0.0

 0.0

Returns:

Type Description
tuple

Tuple containing (grid, averaged_data, smoothed_data)
Source code in toolbox/io/cp2k.py 
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
def get_ave_cube(self, axis=2, gaussian_sigma=0.0):
    """Get averaged cube data along specified axis.

    Parameters
    ----------
    axis : int, optional
        Axis along which to average (0=x, 1=y, 2=z), by default 2
    gaussian_sigma : float, optional
        Standard deviation for Gaussian smoothing, by default 0.0

    Returns
    -------
    tuple
        Tuple containing (grid, averaged_data, smoothed_data)
    """
    if (
        hasattr(self, "axis")
        and self.axis == axis
        and hasattr(self, "ave_cube_data")
    ):
        pass
    else:
        self.axis = axis
        self.ave_grid = self.cube_vectors[self.axis][self.axis] * np.arange(
            self.n_grid[self.axis]
        )
        ave_axis = tuple(np.delete(np.arange(3), self.axis).tolist())
        self.ave_cube_data = np.mean(self.cube_data, axis=ave_axis)

    if gaussian_sigma > 0.0:
        self.ave_cube_data_convolve = gaussian_convolve(
            self.ave_grid, self.ave_cube_data, gaussian_sigma
        )
    else:
        self.ave_cube_data_convolve = copy.deepcopy(self.ave_cube_data)

    return (self.ave_grid, self.ave_cube_data, self.ave_cube_data_convolve)

Cp2kHartreeCube

Bases: Cp2kCube

Class for handling CP2K Hartree potential cube files.

This class extends Cp2kCube to specifically handle Hartree potential data, including methods to calculate potential drops and related properties.

Parameters:

Name Type Description Default
fname str

Path to Hartree potential cube file

 required
vac_region list

Vacuum region boundaries [start, end], by default None

 None
Source code in toolbox/io/cp2k.py 
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
class Cp2kHartreeCube(Cp2kCube):
    """
    Class for handling CP2K Hartree potential cube files.

    This class extends Cp2kCube to specifically handle Hartree potential data,
    including methods to calculate potential drops and related properties.

    Parameters
    ----------
    fname : str
        Path to Hartree potential cube file
    vac_region : list, optional
        Vacuum region boundaries [start, end], by default None
    """

    def __init__(
        self,
        fname,
        vac_region: list = None,
    ) -> None:
        super().__init__(fname)
        if vac_region:
            self.set_vac_region(vac_region)

    def set_vac_region(self, vac_region):
        """Set vacuum region for potential drop calculation.

        Parameters
        ----------
        vac_region : list
            Vacuum region boundaries [start, end]
        """
        assert len(vac_region) == 2
        self.vac_region = vac_region

    def get_ave_cube(self, axis=2, gaussian_sigma=0):
        """Get averaged Hartree potential cube data along specified axis.

        Parameters
        ----------
        axis : int, optional
            Axis along which to average (0=x, 1=y, 2=z), by default 2
        gaussian_sigma : float, optional
            Standard deviation for Gaussian smoothing, by default 0

        Returns
        -------
        tuple
            Tuple containing (grid, averaged_data, smoothed_data) in eV
        """
        if (
            hasattr(self, "axis")
            and self.axis == axis
            and hasattr(self, "ave_cube_data")
        ):
            pass
        else:
            # (self.ave_grid, self.ave_cube_data, self.ave_cube_data_convolve)
            super().get_ave_cube(axis, gaussian_sigma)
            self.ave_cube_data_convolve *= AU_TO_EV
            self.ave_cube_data *= AU_TO_EV
        return (self.ave_grid, self.ave_cube_data, self.ave_cube_data_convolve)

    @property
    def potdrop(self):
        """Calculate potential drop across the vacuum region.

        Returns
        -------
        float
            Potential drop in V
        """
        start_id = np.argmin(np.abs(self.ave_grid - self.vac_region[0]))
        end_id = np.argmin(np.abs(self.ave_grid - self.vac_region[1]))
        if start_id > end_id:
            _data = np.append(
                self.ave_cube_data[start_id:], self.ave_cube_data[:end_id]
            )
        else:
            _data = np.array(self.ave_cube_data)[
                self.vac_region[0] : self.vac_region[1]
            ]
        dev_data = np.diff(_data, axis=0)
        p_jump = np.argmax(np.abs(dev_data))
        return dev_data[p_jump]

    @property
    def dipole(self):
        """Calculate dipole moment from potential drop.

        Returns
        -------
        float
            Dipole moment in e·Å
        """
        d = -self.potdrop * self.cross_area * EPSILON
        return d

    def set_cross_area(self, cross_area):
        """Set cross-sectional area for dipole calculation.

        Parameters
        ----------
        cross_area : float
            Cross-sectional area in Ų
        """
        self.cross_area = cross_area
dipole property

Calculate dipole moment from potential drop.

Returns:

Type Description
float

Dipole moment in e·Å
potdrop property

Calculate potential drop across the vacuum region.

Returns:

Type Description
float

Potential drop in V
get_ave_cube(axis=2, gaussian_sigma=0)

Get averaged Hartree potential cube data along specified axis.

Parameters:

Name Type Description Default
axis int

Axis along which to average (0=x, 1=y, 2=z), by default 2

 2
gaussian_sigma float

Standard deviation for Gaussian smoothing, by default 0

 0

Returns:

Type Description
tuple

Tuple containing (grid, averaged_data, smoothed_data) in eV
Source code in toolbox/io/cp2k.py 
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
def get_ave_cube(self, axis=2, gaussian_sigma=0):
    """Get averaged Hartree potential cube data along specified axis.

    Parameters
    ----------
    axis : int, optional
        Axis along which to average (0=x, 1=y, 2=z), by default 2
    gaussian_sigma : float, optional
        Standard deviation for Gaussian smoothing, by default 0

    Returns
    -------
    tuple
        Tuple containing (grid, averaged_data, smoothed_data) in eV
    """
    if (
        hasattr(self, "axis")
        and self.axis == axis
        and hasattr(self, "ave_cube_data")
    ):
        pass
    else:
        # (self.ave_grid, self.ave_cube_data, self.ave_cube_data_convolve)
        super().get_ave_cube(axis, gaussian_sigma)
        self.ave_cube_data_convolve *= AU_TO_EV
        self.ave_cube_data *= AU_TO_EV
    return (self.ave_grid, self.ave_cube_data, self.ave_cube_data_convolve)
set_cross_area(cross_area)

Set cross-sectional area for dipole calculation.

Parameters:

Name Type Description Default
cross_area float

Cross-sectional area in Ų

 required
Source code in toolbox/io/cp2k.py 
1700
1701
1702
1703
1704
1705
1706
1707
1708
def set_cross_area(self, cross_area):
    """Set cross-sectional area for dipole calculation.

    Parameters
    ----------
    cross_area : float
        Cross-sectional area in Ų
    """
    self.cross_area = cross_area
set_vac_region(vac_region)

Set vacuum region for potential drop calculation.

Parameters:

Name Type Description Default
vac_region list

Vacuum region boundaries [start, end]

 required
Source code in toolbox/io/cp2k.py 
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
def set_vac_region(self, vac_region):
    """Set vacuum region for potential drop calculation.

    Parameters
    ----------
    vac_region : list
        Vacuum region boundaries [start, end]
    """
    assert len(vac_region) == 2
    self.vac_region = vac_region

Cp2kInput

Class for CP2K input file generation (on the basis of templates).

Attributes:

Name Type Description
atoms ASE Atoms object

TBC
input_type str

TBC
pp_dir str

directory for basis set, peusudopotential, etc.
wfn_restart str

wfn file for restart, see ref:
qm_charge float

charge in QS
multiplicity int

ref:
uks boolen

ref:
cutoff int

ref:
rel_cutoff int

ref:

Examples:

>>> from ase import io
>>> from toolbox.io.cp2k import Cp2kInput
>>> atoms = io.read("POSCAR")
>>> input = Cp2kInput(atoms,
>>>                   pp_dir="/data/basis",
>>>                   hartree=True,
>>>                   eden=True)
>>> input.write()
Source code in toolbox/io/cp2k.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
 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
112
113
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
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
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
300
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
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
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
421
422
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
466
467
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
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
575
576
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
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
class Cp2kInput:
    """
    Class for CP2K input file generation (on the basis of templates).

    Attributes
    ----------
    atoms: ASE Atoms object
        TBC
    input_type: str
        TBC
    pp_dir: str
        directory for basis set, peusudopotential, etc.
    wfn_restart: str
        wfn file for restart, see ref:
    qm_charge: float
        charge in QS
    multiplicity: int
        ref:
    uks: boolen
        ref:
    cutoff: int
        ref:
    rel_cutoff: int
        ref:

    Examples
    --------
    >>> from ase import io
    >>> from toolbox.io.cp2k import Cp2kInput
    >>> atoms = io.read("POSCAR")
    >>> input = Cp2kInput(atoms,
    >>>                   pp_dir="/data/basis",
    >>>                   hartree=True,
    >>>                   eden=True)
    >>> input.write()
    """

    def __init__(self, atoms, input_type="energy", **kwargs) -> None:
        self.atoms = atoms
        self.input_dict = copy.deepcopy(cp2k_default_input[input_type])
        # print(kwargs)
        # read user setup in config file
        try:
            update_d = CONFIGS["io"]["cp2k"]["input"]
        except KeyError:
            update_d = {}
        update_dict(kwargs, update_d)
        self.set_params(kwargs)

    def set_params(self, kwargs):
        """Set parameters for CP2K input.

        Parameters
        ----------
        kwargs : dict
            Dictionary of parameters to set
        """
        for kw, value in kwargs.items():
            update_d = getattr(self, f"set_{kw}")(value)
            update_dict(self.input_dict, update_d)

    def write(self, output_dir=".", fp_params=None, save_dict=False):
        """
        Generate coord.xyz and input.inp for CP2K calculation at output_dir.

        Parameters
        ----------
        output_dir : str
            directory to store coord.xyz and input.inp
        fp_params : dict
            dict for updated parameters
        """
        if fp_params is None:
            fp_params = {}
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        cell = self.atoms.get_cell()
        cell_a = np.array2string(
            cell[0], formatter={"float_kind": lambda x: f"{x:.4f}"}
        )
        cell_a = cell_a[1:-1]
        cell_b = np.array2string(
            cell[1], formatter={"float_kind": lambda x: f"{x:.4f}"}
        )
        cell_b = cell_b[1:-1]
        cell_c = np.array2string(
            cell[2], formatter={"float_kind": lambda x: f"{x:.4f}"}
        )
        cell_c = cell_c[1:-1]

        user_config = fp_params
        update_dict(self.input_dict, user_config)

        if self.input_dict["FORCE_EVAL"].get("QMMM", None) is not None:
            cell_config = {
                "FORCE_EVAL": {
                    "SUBSYS": {"CELL": {"A": cell_a, "B": cell_b, "C": cell_c}},
                    "QMMM": {
                        "CELL": {
                            "A": cell_a,
                            "B": cell_b,
                            "C": cell_c,
                            "PERIODIC": "XYZ",
                        }
                    },
                }
            }
        else:
            cell_config = {
                "FORCE_EVAL": {
                    "SUBSYS": {"CELL": {"A": cell_a, "B": cell_b, "C": cell_c}}
                }
            }

        update_dict(self.input_dict, cell_config)
        # output list
        input_str = iterdict(self.input_dict, out_list=["\n"], loop_idx=0)
        # del input_str[0]
        # del input_str[-1]
        # print(input_str)
        str = "\n".join(input_str)
        str = str.strip("\n")

        io.write(os.path.join(output_dir, "coord.xyz"), self.atoms)
        with open(os.path.join(output_dir, "input.inp"), "w", encoding="utf-8") as f:
            f.write(str)

        if save_dict:
            save_dict_json(self.input_dict, os.path.join(output_dir, "input.json"))

    def set_project(self, project_name: str):
        """Set project name for CP2K calculation.

        Parameters
        ----------
        project_name : str
            Project name for CP2K calculation

        Returns
        -------
        dict
            Update dictionary for project name
        """
        update_d = {"GLOBAL": {"PROJECT": project_name}}
        return update_d

    def set_pp_dir(self, pp_dir):
        """Set pseudopotential directory.

        Parameters
        ----------
        pp_dir : str
            Directory containing basis sets and pseudopotentials

        Returns
        -------
        dict
            Update dictionary for pseudopotential directory
        """
        pp_dir = os.path.abspath(pp_dir)
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "BASIS_SET_FILE_NAME": [
                        os.path.join(pp_dir, "BASIS_MOLOPT"),
                        os.path.join(pp_dir, "BASIS_ADMM"),
                        os.path.join(pp_dir, "BASIS_ADMM_MOLOPT"),
                        os.path.join(pp_dir, "BASIS_MOLOPT-HSE06"),
                    ],
                    "POTENTIAL_FILE_NAME": os.path.join(pp_dir, "GTH_POTENTIALS"),
                    "XC": {
                        "vdW_POTENTIAL": {
                            "PAIR_POTENTIAL": {
                                "PARAMETER_FILE_NAME": os.path.join(pp_dir, "dftd3.dat")
                            }
                        }
                    },
                }
            }
        }
        return update_d

    def set_wfn_restart(self, wfn_file):
        """Set wavefunction restart file.

        Parameters
        ----------
        wfn_file : str or None
            Path to wavefunction restart file

        Returns
        -------
        dict
            Update dictionary for wavefunction restart file
        """
        update_d = {}
        if wfn_file is not None:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {"WFN_RESTART_FILE_NAME": os.path.abspath(wfn_file)}
                }
            }
        return update_d

    def set_qm_charge(self, charge):
        """Set quantum mechanical charge.

        Parameters
        ----------
        charge : float
            Total charge of the system

        Returns
        -------
        dict
            Update dictionary for charge
        """
        update_d = {"FORCE_EVAL": {"DFT": {"CHARGE": charge}}}
        return update_d

    def set_multiplicity(self, multiplicity):
        """Set spin multiplicity.

        Parameters
        ----------
        multiplicity : int
            Spin multiplicity of the system

        Returns
        -------
        dict
            Update dictionary for multiplicity
        """
        update_d = {"FORCE_EVAL": {"DFT": {"MULTIPLICITY": multiplicity}}}
        return update_d

    def set_uks(self, flag):
        """Set unrestricted Kohn-Sham calculation.

        Parameters
        ----------
        flag : bool
            Whether to use UKS calculation

        Returns
        -------
        dict
            Update dictionary for UKS
        """
        if flag:
            update_d = {"FORCE_EVAL": {"DFT": {"UKS": ".TRUE."}}}
            return update_d
        else:
            return {}

    def set_cutoff(self, cutoff):
        """Set plane wave cutoff.

        Parameters
        ----------
        cutoff : float
            Plane wave cutoff in Ry

        Returns
        -------
        dict
            Update dictionary for cutoff
        """
        update_d = {"FORCE_EVAL": {"DFT": {"MGRID": {"CUTOFF": cutoff}}}}
        return update_d

    def set_rel_cutoff(self, rel_cutoff):
        """Set relative cutoff.

        Parameters
        ----------
        rel_cutoff : float
            Relative cutoff for multi-grid

        Returns
        -------
        dict
            Update dictionary for relative cutoff
        """
        update_d = {"FORCE_EVAL": {"DFT": {"MGRID": {"REL_CUTOFF": rel_cutoff}}}}
        return update_d

    def set_kp(self, kp_mp):
        """Set k-points mesh.

        Parameters
        ----------
        kp_mp : tuple
            K-point mesh as (kx, ky, kz)

        Returns
        -------
        dict
            Update dictionary for k-points
        """
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "KPOINTS": {
                        "SCHEME MONKHORST-PACK": f"{kp_mp[0]:d} {kp_mp[1]:d} {kp_mp[2]:d}",
                        "SYMMETRY": ".TRUE.",
                        "EPS_GEO": 1.0e-8,
                        "FULL_GRID": ".TRUE.",
                        "PARALLEL_GROUP_SIZE": 0,
                    }
                }
            }
        }
        return update_d

    def set_max_scf(self, max_scf: int):
        """Set maximum SCF iterations.

        Parameters
        ----------
        max_scf : int
            Maximum number of SCF iterations

        Returns
        -------
        dict
            Update dictionary for maximum SCF iterations
        """
        update_d = {"FORCE_EVAL": {"DFT": {"SCF": {"MAX_SCF": max_scf}}}}
        return update_d

    def set_eps_scf(self, eps_scf: float):
        """Set SCF convergence threshold.

        Parameters
        ----------
        eps_scf : float
            SCF convergence threshold

        Returns
        -------
        dict
            Update dictionary for SCF convergence threshold
        """
        update_d = {"FORCE_EVAL": {"DFT": {"SCF": {"EPS_SCF": eps_scf}}}}
        return update_d

    def set_dip_cor(self, flag):
        """Set surface dipole correction.

        Parameters
        ----------
        flag : bool
            Whether to apply surface dipole correction

        Returns
        -------
        dict
            Update dictionary for dipole correction
        """
        if flag:
            update_d = {"FORCE_EVAL": {"DFT": {"SURFACE_DIPOLE_CORRECTION": ".TRUE."}}}
            return update_d
        else:
            return {}

    def set_eden(self, flag):
        """Set electron density cube output.

        Parameters
        ----------
        flag : bool
            Whether to output electron density cube

        Returns
        -------
        dict
            Update dictionary for electron density output
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "PRINT": {
                            "E_DENSITY_CUBE": {"ADD_LAST": "NUMERIC", "STRIDE": "8 8 1"}
                        }
                    }
                }
            }
            return update_d
        else:
            return {}

    def set_mo(self, flag):
        """Set molecular orbital cube output.

        Parameters
        ----------
        flag : bool
            Whether to output molecular orbital cubes

        Returns
        -------
        dict
            Update dictionary for MO output
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {"DFT": {"PRINT": {"MO_CUBES": {"ADD_LAST": "NUMERIC"}}}}
            }
            return update_d
        else:
            return {}

    def set_pdos(self, flag):
        """Set projected density of states output.

        Parameters
        ----------
        flag : bool
            Whether to output PDOS

        Returns
        -------
        dict
            Update dictionary for PDOS output
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "PRINT": {
                            "PDOS": {
                                "COMPONENTS": ".TRUE.",
                                "ADD_LAST": "NUMERIC",
                                "NLUMO": -1,
                                "COMMON_ITERATION_LEVELS": 0,
                            }
                        }
                    }
                }
            }
            return update_d
        else:
            return {}

    def set_hartree(self, flag):
        """Set Hartree potential cube output.

        Parameters
        ----------
        flag : bool
            Whether to output Hartree potential cube

        Returns
        -------
        dict
            Update dictionary for Hartree potential output
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "PRINT": {
                            "V_HARTREE_CUBE": {"ADD_LAST": "NUMERIC", "STRIDE": "8 8 1"}
                        }
                    }
                }
            }
            return update_d
        else:
            return {}

    def set_efield(self, flag):
        """Set electric field cube output.

        Parameters
        ----------
        flag : bool
            Whether to output electric field cube

        Returns
        -------
        dict
            Update dictionary for electric field output
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "PRINT": {
                            "EFIELD_CUBE": {"ADD_LAST": "NUMERIC", "STRIDE": "8 8 1"}
                        }
                    }
                }
            }
            return update_d
        else:
            return {}

    def set_totden(self, flag):
        """Set total density cube output.

        Parameters
        ----------
        flag : bool
            Whether to output total density cube

        Returns
        -------
        dict
            Update dictionary for total density output
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "PRINT": {
                            "TOT_DENSITY_CUBE": {
                                "ADD_LAST": "NUMERIC",
                                "STRIDE": "1 1 1",
                            }
                        }
                    }
                }
            }
            return update_d
        else:
            return {}

    def set_extended_fft_lengths(self, flag):
        """Set extended FFT lengths.

        Parameters
        ----------
        flag : bool
            Whether to use extended FFT lengths

        Returns
        -------
        dict
            Update dictionary for extended FFT lengths
        """
        if flag:
            update_d = {"GLOBAL": {"EXTENDED_FFT_LENGTHS": ".TRUE."}}
            return update_d
        else:
            return {}

    def set_smear(self, flag):
        """Set smearing method for SCF calculation.

        Parameters
        ----------
        flag : bool
            Whether to enable smearing

        Returns
        -------
        dict
            Update dictionary for smearing settings
        """
        if not flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "SCF": {
                            "ADDED_MOS": 0,
                            "CHOLESKY": "RESTORE",
                            "SMEAR": {"_": ".FALSE."},
                            "DIAGONALIZATION": {"_": ".FALSE."},
                        }
                    }
                }
            }
            return update_d

    def set_mlwf(self, flag):
        """Set maximally localized Wannier functions.

        Parameters
        ----------
        flag : bool
            Whether to enable MLWF calculation

        Returns
        -------
        dict
            Update dictionary for MLWF settings
        """
        if flag:
            update_d = {
                "FORCE_EVAL": {
                    "DFT": {
                        "LOCALIZE": {
                            "METHOD": "CRAZY",
                            "EPS_LOCALIZATION": 1e-08,
                            "PRINT": {"WANNIER_CENTERS": {"IONS+CENTERS": ".TRUE."}},
                        }
                    }
                }
            }
            return update_d
        else:
            return {}

    def set_kind(self, kind_dict: dict):
        """
        Set atom kind parameters.

        Parameters
        ----------
        kind_dict : dict
            dict to update kind section, for example:
            {
                "S": {
                    "ELEMENT": "O",
                    "BASIS_SET": "DZVP-MOLOPT-SR-GTH",
                    "POTENTIAL": "GTH-PBE-q6"
                },
                "Li": {
                    "ELEMENT": "H",
                    "BASIS_SET": "DZVP-MOLOPT-SR-GTH",
                    "POTENTIAL": "GTH-PBE-q1"
                }
            }
        """
        update_d = {"FORCE_EVAL": {"SUBSYS": {}}}
        update_dict(self.input_dict, update_d)

        old_kind_list = self.input_dict["FORCE_EVAL"]["SUBSYS"].get("KIND", [])
        if len(old_kind_list) > 0:
            for k, v in kind_dict.items():
                tmp_dict = copy.deepcopy(v)
                tmp_dict.update({"_": k})
                flag = False
                for ii, item in enumerate(old_kind_list):
                    if k == item["_"]:
                        # print(v)
                        old_kind_list[ii] = tmp_dict
                        flag = True
                        break
                if not flag:
                    old_kind_list.append(tmp_dict)
        return {}

    def set_restart(self, flag):
        """Set restart file for CP2K calculation.

        Parameters
        ----------
        flag : bool
            Whether to use restart file

        Returns
        -------
        dict
            Update dictionary for restart settings
        """
        if flag:
            update_d = {
                "EXT_RESTART": {
                    "RESTART_FILE_NAME": "{}-1.restart".format(
                        self.input_dict["GLOBAL"]["PROJECT"]
                    )
                }
            }
            return update_d
        else:
            return {}

    def set_md_step(self, md_step):
        """Set number of MD steps.

        Parameters
        ----------
        md_step : int
            Number of MD steps

        Returns
        -------
        dict
            Update dictionary for MD steps
        """
        update_d = {"MOTION": {"MD": {"STEPS": md_step}}}
        return update_d

    def set_md_temp(self, md_temp):
        """Set MD temperature.

        Parameters
        ----------
        md_temp : float
            Temperature for MD simulation

        Returns
        -------
        dict
            Update dictionary for MD temperature
        """
        update_d = {"MOTION": {"MD": {"TEMPERATURE": md_temp}}}
        return update_d

    def set_md_timestep(self, md_timestep):
        """Set MD timestep.

        Parameters
        ----------
        md_timestep : float
            Timestep for MD simulation

        Returns
        -------
        dict
            Update dictionary for MD timestep
        """
        update_d = {"MOTION": {"MD": {"TIMESTEP": md_timestep}}}
        return update_d

    def update_dp(self, dp_model: str):
        """Update CP2K input with Deep Potential model.

        Parameters
        ----------
        dp_model : str
            Path to Deep Potential model file
        """
        from deepmd.infer import DeepPot

        dp = DeepPot(dp_model)
        type_map = dp.tmap
        for ii, atype in enumerate(type_map):
            self.input_dict["FORCE_EVAL"]["MM"]["FORCEFIELD"]["CHARGE"].append(
                {
                    "ATOM": atype,
                    "CHARGE": 0.0,
                }
            )
            self.input_dict["FORCE_EVAL"]["MM"]["FORCEFIELD"]["NONBONDED"][
                "DEEPMD"
            ].append(
                {
                    "ATOMS": f"{atype} {atype}",
                    "POT_FILE_NAME": dp_model,
                    "ATOM_DEEPMD_TYPE": ii,
                }
            )
set_cutoff(cutoff)

Set plane wave cutoff.

Parameters:

Name Type Description Default
cutoff float

Plane wave cutoff in Ry

 required

Returns:

Type Description
dict

Update dictionary for cutoff
Source code in toolbox/io/cp2k.py 
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def set_cutoff(self, cutoff):
    """Set plane wave cutoff.

    Parameters
    ----------
    cutoff : float
        Plane wave cutoff in Ry

    Returns
    -------
    dict
        Update dictionary for cutoff
    """
    update_d = {"FORCE_EVAL": {"DFT": {"MGRID": {"CUTOFF": cutoff}}}}
    return update_d
set_dip_cor(flag)

Set surface dipole correction.

Parameters:

Name Type Description Default
flag bool

Whether to apply surface dipole correction

 required

Returns:

Type Description
dict

Update dictionary for dipole correction
Source code in toolbox/io/cp2k.py 
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
def set_dip_cor(self, flag):
    """Set surface dipole correction.

    Parameters
    ----------
    flag : bool
        Whether to apply surface dipole correction

    Returns
    -------
    dict
        Update dictionary for dipole correction
    """
    if flag:
        update_d = {"FORCE_EVAL": {"DFT": {"SURFACE_DIPOLE_CORRECTION": ".TRUE."}}}
        return update_d
    else:
        return {}
set_eden(flag)

Set electron density cube output.

Parameters:

Name Type Description Default
flag bool

Whether to output electron density cube

 required

Returns:

Type Description
dict

Update dictionary for electron density output
Source code in toolbox/io/cp2k.py 
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 set_eden(self, flag):
    """Set electron density cube output.

    Parameters
    ----------
    flag : bool
        Whether to output electron density cube

    Returns
    -------
    dict
        Update dictionary for electron density output
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "PRINT": {
                        "E_DENSITY_CUBE": {"ADD_LAST": "NUMERIC", "STRIDE": "8 8 1"}
                    }
                }
            }
        }
        return update_d
    else:
        return {}
set_efield(flag)

Set electric field cube output.

Parameters:

Name Type Description Default
flag bool

Whether to output electric field cube

 required

Returns:

Type Description
dict

Update dictionary for electric field output
Source code in toolbox/io/cp2k.py 
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
def set_efield(self, flag):
    """Set electric field cube output.

    Parameters
    ----------
    flag : bool
        Whether to output electric field cube

    Returns
    -------
    dict
        Update dictionary for electric field output
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "PRINT": {
                        "EFIELD_CUBE": {"ADD_LAST": "NUMERIC", "STRIDE": "8 8 1"}
                    }
                }
            }
        }
        return update_d
    else:
        return {}
set_eps_scf(eps_scf)

Set SCF convergence threshold.

Parameters:

Name Type Description Default
eps_scf float

SCF convergence threshold

 required

Returns:

Type Description
dict

Update dictionary for SCF convergence threshold
Source code in toolbox/io/cp2k.py 
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
def set_eps_scf(self, eps_scf: float):
    """Set SCF convergence threshold.

    Parameters
    ----------
    eps_scf : float
        SCF convergence threshold

    Returns
    -------
    dict
        Update dictionary for SCF convergence threshold
    """
    update_d = {"FORCE_EVAL": {"DFT": {"SCF": {"EPS_SCF": eps_scf}}}}
    return update_d
set_extended_fft_lengths(flag)

Set extended FFT lengths.

Parameters:

Name Type Description Default
flag bool

Whether to use extended FFT lengths

 required

Returns:

Type Description
dict

Update dictionary for extended FFT lengths
Source code in toolbox/io/cp2k.py 
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
def set_extended_fft_lengths(self, flag):
    """Set extended FFT lengths.

    Parameters
    ----------
    flag : bool
        Whether to use extended FFT lengths

    Returns
    -------
    dict
        Update dictionary for extended FFT lengths
    """
    if flag:
        update_d = {"GLOBAL": {"EXTENDED_FFT_LENGTHS": ".TRUE."}}
        return update_d
    else:
        return {}
set_hartree(flag)

Set Hartree potential cube output.

Parameters:

Name Type Description Default
flag bool

Whether to output Hartree potential cube

 required

Returns:

Type Description
dict

Update dictionary for Hartree potential output
Source code in toolbox/io/cp2k.py 
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
500
def set_hartree(self, flag):
    """Set Hartree potential cube output.

    Parameters
    ----------
    flag : bool
        Whether to output Hartree potential cube

    Returns
    -------
    dict
        Update dictionary for Hartree potential output
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "PRINT": {
                        "V_HARTREE_CUBE": {"ADD_LAST": "NUMERIC", "STRIDE": "8 8 1"}
                    }
                }
            }
        }
        return update_d
    else:
        return {}
set_kind(kind_dict)

Set atom kind parameters.

Parameters:

Name Type Description Default
kind_dict dict

dict to update kind section, for example: { "S": { "ELEMENT": "O", "BASIS_SET": "DZVP-MOLOPT-SR-GTH", "POTENTIAL": "GTH-PBE-q6" }, "Li": { "ELEMENT": "H", "BASIS_SET": "DZVP-MOLOPT-SR-GTH", "POTENTIAL": "GTH-PBE-q1" } }

 required
Source code in toolbox/io/cp2k.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
def set_kind(self, kind_dict: dict):
    """
    Set atom kind parameters.

    Parameters
    ----------
    kind_dict : dict
        dict to update kind section, for example:
        {
            "S": {
                "ELEMENT": "O",
                "BASIS_SET": "DZVP-MOLOPT-SR-GTH",
                "POTENTIAL": "GTH-PBE-q6"
            },
            "Li": {
                "ELEMENT": "H",
                "BASIS_SET": "DZVP-MOLOPT-SR-GTH",
                "POTENTIAL": "GTH-PBE-q1"
            }
        }
    """
    update_d = {"FORCE_EVAL": {"SUBSYS": {}}}
    update_dict(self.input_dict, update_d)

    old_kind_list = self.input_dict["FORCE_EVAL"]["SUBSYS"].get("KIND", [])
    if len(old_kind_list) > 0:
        for k, v in kind_dict.items():
            tmp_dict = copy.deepcopy(v)
            tmp_dict.update({"_": k})
            flag = False
            for ii, item in enumerate(old_kind_list):
                if k == item["_"]:
                    # print(v)
                    old_kind_list[ii] = tmp_dict
                    flag = True
                    break
            if not flag:
                old_kind_list.append(tmp_dict)
    return {}
set_kp(kp_mp)

Set k-points mesh.

Parameters:

Name Type Description Default
kp_mp tuple

K-point mesh as (kx, ky, kz)

 required

Returns:

Type Description
dict

Update dictionary for k-points
Source code in toolbox/io/cp2k.py 
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
def set_kp(self, kp_mp):
    """Set k-points mesh.

    Parameters
    ----------
    kp_mp : tuple
        K-point mesh as (kx, ky, kz)

    Returns
    -------
    dict
        Update dictionary for k-points
    """
    update_d = {
        "FORCE_EVAL": {
            "DFT": {
                "KPOINTS": {
                    "SCHEME MONKHORST-PACK": f"{kp_mp[0]:d} {kp_mp[1]:d} {kp_mp[2]:d}",
                    "SYMMETRY": ".TRUE.",
                    "EPS_GEO": 1.0e-8,
                    "FULL_GRID": ".TRUE.",
                    "PARALLEL_GROUP_SIZE": 0,
                }
            }
        }
    }
    return update_d
set_max_scf(max_scf)

Set maximum SCF iterations.

Parameters:

Name Type Description Default
max_scf int

Maximum number of SCF iterations

 required

Returns:

Type Description
dict

Update dictionary for maximum SCF iterations
Source code in toolbox/io/cp2k.py 
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
def set_max_scf(self, max_scf: int):
    """Set maximum SCF iterations.

    Parameters
    ----------
    max_scf : int
        Maximum number of SCF iterations

    Returns
    -------
    dict
        Update dictionary for maximum SCF iterations
    """
    update_d = {"FORCE_EVAL": {"DFT": {"SCF": {"MAX_SCF": max_scf}}}}
    return update_d
set_md_step(md_step)

Set number of MD steps.

Parameters:

Name Type Description Default
md_step int

Number of MD steps

 required

Returns:

Type Description
dict

Update dictionary for MD steps
Source code in toolbox/io/cp2k.py 
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
def set_md_step(self, md_step):
    """Set number of MD steps.

    Parameters
    ----------
    md_step : int
        Number of MD steps

    Returns
    -------
    dict
        Update dictionary for MD steps
    """
    update_d = {"MOTION": {"MD": {"STEPS": md_step}}}
    return update_d
set_md_temp(md_temp)

Set MD temperature.

Parameters:

Name Type Description Default
md_temp float

Temperature for MD simulation

 required

Returns:

Type Description
dict

Update dictionary for MD temperature
Source code in toolbox/io/cp2k.py 
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
def set_md_temp(self, md_temp):
    """Set MD temperature.

    Parameters
    ----------
    md_temp : float
        Temperature for MD simulation

    Returns
    -------
    dict
        Update dictionary for MD temperature
    """
    update_d = {"MOTION": {"MD": {"TEMPERATURE": md_temp}}}
    return update_d
set_md_timestep(md_timestep)

Set MD timestep.

Parameters:

Name Type Description Default
md_timestep float

Timestep for MD simulation

 required

Returns:

Type Description
dict

Update dictionary for MD timestep
Source code in toolbox/io/cp2k.py 
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
def set_md_timestep(self, md_timestep):
    """Set MD timestep.

    Parameters
    ----------
    md_timestep : float
        Timestep for MD simulation

    Returns
    -------
    dict
        Update dictionary for MD timestep
    """
    update_d = {"MOTION": {"MD": {"TIMESTEP": md_timestep}}}
    return update_d
set_mlwf(flag)

Set maximally localized Wannier functions.

Parameters:

Name Type Description Default
flag bool

Whether to enable MLWF calculation

 required

Returns:

Type Description
dict

Update dictionary for MLWF settings
Source code in toolbox/io/cp2k.py 
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 set_mlwf(self, flag):
    """Set maximally localized Wannier functions.

    Parameters
    ----------
    flag : bool
        Whether to enable MLWF calculation

    Returns
    -------
    dict
        Update dictionary for MLWF settings
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "LOCALIZE": {
                        "METHOD": "CRAZY",
                        "EPS_LOCALIZATION": 1e-08,
                        "PRINT": {"WANNIER_CENTERS": {"IONS+CENTERS": ".TRUE."}},
                    }
                }
            }
        }
        return update_d
    else:
        return {}
set_mo(flag)

Set molecular orbital cube output.

Parameters:

Name Type Description Default
flag bool

Whether to output molecular orbital cubes

 required

Returns:

Type Description
dict

Update dictionary for MO output
Source code in toolbox/io/cp2k.py 
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
def set_mo(self, flag):
    """Set molecular orbital cube output.

    Parameters
    ----------
    flag : bool
        Whether to output molecular orbital cubes

    Returns
    -------
    dict
        Update dictionary for MO output
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {"DFT": {"PRINT": {"MO_CUBES": {"ADD_LAST": "NUMERIC"}}}}
        }
        return update_d
    else:
        return {}
set_multiplicity(multiplicity)

Set spin multiplicity.

Parameters:

Name Type Description Default
multiplicity int

Spin multiplicity of the system

 required

Returns:

Type Description
dict

Update dictionary for multiplicity
Source code in toolbox/io/cp2k.py 
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def set_multiplicity(self, multiplicity):
    """Set spin multiplicity.

    Parameters
    ----------
    multiplicity : int
        Spin multiplicity of the system

    Returns
    -------
    dict
        Update dictionary for multiplicity
    """
    update_d = {"FORCE_EVAL": {"DFT": {"MULTIPLICITY": multiplicity}}}
    return update_d
set_params(kwargs)

Set parameters for CP2K input.

Parameters:

Name Type Description Default
kwargs dict

Dictionary of parameters to set

 required
Source code in toolbox/io/cp2k.py 
77
78
79
80
81
82
83
84
85
86
87
def set_params(self, kwargs):
    """Set parameters for CP2K input.

    Parameters
    ----------
    kwargs : dict
        Dictionary of parameters to set
    """
    for kw, value in kwargs.items():
        update_d = getattr(self, f"set_{kw}")(value)
        update_dict(self.input_dict, update_d)
set_pdos(flag)

Set projected density of states output.

Parameters:

Name Type Description Default
flag bool

Whether to output PDOS

 required

Returns:

Type Description
dict

Update dictionary for PDOS output
Source code in toolbox/io/cp2k.py 
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
def set_pdos(self, flag):
    """Set projected density of states output.

    Parameters
    ----------
    flag : bool
        Whether to output PDOS

    Returns
    -------
    dict
        Update dictionary for PDOS output
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "PRINT": {
                        "PDOS": {
                            "COMPONENTS": ".TRUE.",
                            "ADD_LAST": "NUMERIC",
                            "NLUMO": -1,
                            "COMMON_ITERATION_LEVELS": 0,
                        }
                    }
                }
            }
        }
        return update_d
    else:
        return {}
set_pp_dir(pp_dir)

Set pseudopotential directory.

Parameters:

Name Type Description Default
pp_dir str

Directory containing basis sets and pseudopotentials

 required

Returns:

Type Description
dict

Update dictionary for pseudopotential directory
Source code in toolbox/io/cp2k.py 
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
def set_pp_dir(self, pp_dir):
    """Set pseudopotential directory.

    Parameters
    ----------
    pp_dir : str
        Directory containing basis sets and pseudopotentials

    Returns
    -------
    dict
        Update dictionary for pseudopotential directory
    """
    pp_dir = os.path.abspath(pp_dir)
    update_d = {
        "FORCE_EVAL": {
            "DFT": {
                "BASIS_SET_FILE_NAME": [
                    os.path.join(pp_dir, "BASIS_MOLOPT"),
                    os.path.join(pp_dir, "BASIS_ADMM"),
                    os.path.join(pp_dir, "BASIS_ADMM_MOLOPT"),
                    os.path.join(pp_dir, "BASIS_MOLOPT-HSE06"),
                ],
                "POTENTIAL_FILE_NAME": os.path.join(pp_dir, "GTH_POTENTIALS"),
                "XC": {
                    "vdW_POTENTIAL": {
                        "PAIR_POTENTIAL": {
                            "PARAMETER_FILE_NAME": os.path.join(pp_dir, "dftd3.dat")
                        }
                    }
                },
            }
        }
    }
    return update_d
set_project(project_name)

Set project name for CP2K calculation.

Parameters:

Name Type Description Default
project_name str

Project name for CP2K calculation

 required

Returns:

Type Description
dict

Update dictionary for project name
Source code in toolbox/io/cp2k.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def set_project(self, project_name: str):
    """Set project name for CP2K calculation.

    Parameters
    ----------
    project_name : str
        Project name for CP2K calculation

    Returns
    -------
    dict
        Update dictionary for project name
    """
    update_d = {"GLOBAL": {"PROJECT": project_name}}
    return update_d
set_qm_charge(charge)

Set quantum mechanical charge.

Parameters:

Name Type Description Default
charge float

Total charge of the system

 required

Returns:

Type Description
dict

Update dictionary for charge
Source code in toolbox/io/cp2k.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
def set_qm_charge(self, charge):
    """Set quantum mechanical charge.

    Parameters
    ----------
    charge : float
        Total charge of the system

    Returns
    -------
    dict
        Update dictionary for charge
    """
    update_d = {"FORCE_EVAL": {"DFT": {"CHARGE": charge}}}
    return update_d
set_rel_cutoff(rel_cutoff)

Set relative cutoff.

Parameters:

Name Type Description Default
rel_cutoff float

Relative cutoff for multi-grid

 required

Returns:

Type Description
dict

Update dictionary for relative cutoff
Source code in toolbox/io/cp2k.py 
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
def set_rel_cutoff(self, rel_cutoff):
    """Set relative cutoff.

    Parameters
    ----------
    rel_cutoff : float
        Relative cutoff for multi-grid

    Returns
    -------
    dict
        Update dictionary for relative cutoff
    """
    update_d = {"FORCE_EVAL": {"DFT": {"MGRID": {"REL_CUTOFF": rel_cutoff}}}}
    return update_d
set_restart(flag)

Set restart file for CP2K calculation.

Parameters:

Name Type Description Default
flag bool

Whether to use restart file

 required

Returns:

Type Description
dict

Update dictionary for restart settings
Source code in toolbox/io/cp2k.py 
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
def set_restart(self, flag):
    """Set restart file for CP2K calculation.

    Parameters
    ----------
    flag : bool
        Whether to use restart file

    Returns
    -------
    dict
        Update dictionary for restart settings
    """
    if flag:
        update_d = {
            "EXT_RESTART": {
                "RESTART_FILE_NAME": "{}-1.restart".format(
                    self.input_dict["GLOBAL"]["PROJECT"]
                )
            }
        }
        return update_d
    else:
        return {}
set_smear(flag)

Set smearing method for SCF calculation.

Parameters:

Name Type Description Default
flag bool

Whether to enable smearing

 required

Returns:

Type Description
dict

Update dictionary for smearing settings
Source code in toolbox/io/cp2k.py 
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
def set_smear(self, flag):
    """Set smearing method for SCF calculation.

    Parameters
    ----------
    flag : bool
        Whether to enable smearing

    Returns
    -------
    dict
        Update dictionary for smearing settings
    """
    if not flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "SCF": {
                        "ADDED_MOS": 0,
                        "CHOLESKY": "RESTORE",
                        "SMEAR": {"_": ".FALSE."},
                        "DIAGONALIZATION": {"_": ".FALSE."},
                    }
                }
            }
        }
        return update_d
set_totden(flag)

Set total density cube output.

Parameters:

Name Type Description Default
flag bool

Whether to output total density cube

 required

Returns:

Type Description
dict

Update dictionary for total density output
Source code in toolbox/io/cp2k.py 
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
def set_totden(self, flag):
    """Set total density cube output.

    Parameters
    ----------
    flag : bool
        Whether to output total density cube

    Returns
    -------
    dict
        Update dictionary for total density output
    """
    if flag:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {
                    "PRINT": {
                        "TOT_DENSITY_CUBE": {
                            "ADD_LAST": "NUMERIC",
                            "STRIDE": "1 1 1",
                        }
                    }
                }
            }
        }
        return update_d
    else:
        return {}
set_uks(flag)

Set unrestricted Kohn-Sham calculation.

Parameters:

Name Type Description Default
flag bool

Whether to use UKS calculation

 required

Returns:

Type Description
dict

Update dictionary for UKS
Source code in toolbox/io/cp2k.py 
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def set_uks(self, flag):
    """Set unrestricted Kohn-Sham calculation.

    Parameters
    ----------
    flag : bool
        Whether to use UKS calculation

    Returns
    -------
    dict
        Update dictionary for UKS
    """
    if flag:
        update_d = {"FORCE_EVAL": {"DFT": {"UKS": ".TRUE."}}}
        return update_d
    else:
        return {}
set_wfn_restart(wfn_file)

Set wavefunction restart file.

Parameters:

Name Type Description Default
wfn_file str or None

Path to wavefunction restart file

 required

Returns:

Type Description
dict

Update dictionary for wavefunction restart file
Source code in toolbox/io/cp2k.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
def set_wfn_restart(self, wfn_file):
    """Set wavefunction restart file.

    Parameters
    ----------
    wfn_file : str or None
        Path to wavefunction restart file

    Returns
    -------
    dict
        Update dictionary for wavefunction restart file
    """
    update_d = {}
    if wfn_file is not None:
        update_d = {
            "FORCE_EVAL": {
                "DFT": {"WFN_RESTART_FILE_NAME": os.path.abspath(wfn_file)}
            }
        }
    return update_d
update_dp(dp_model)

Update CP2K input with Deep Potential model.

Parameters:

Name Type Description Default
dp_model str

Path to Deep Potential model file

 required
Source code in toolbox/io/cp2k.py 
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
def update_dp(self, dp_model: str):
    """Update CP2K input with Deep Potential model.

    Parameters
    ----------
    dp_model : str
        Path to Deep Potential model file
    """
    from deepmd.infer import DeepPot

    dp = DeepPot(dp_model)
    type_map = dp.tmap
    for ii, atype in enumerate(type_map):
        self.input_dict["FORCE_EVAL"]["MM"]["FORCEFIELD"]["CHARGE"].append(
            {
                "ATOM": atype,
                "CHARGE": 0.0,
            }
        )
        self.input_dict["FORCE_EVAL"]["MM"]["FORCEFIELD"]["NONBONDED"][
            "DEEPMD"
        ].append(
            {
                "ATOMS": f"{atype} {atype}",
                "POT_FILE_NAME": dp_model,
                "ATOM_DEEPMD_TYPE": ii,
            }
        )
write(output_dir='.', fp_params=None, save_dict=False)

Generate coord.xyz and input.inp for CP2K calculation at output_dir.

Parameters:

Name Type Description Default
output_dir str

directory to store coord.xyz and input.inp

 '.'
fp_params dict

dict for updated parameters

 None
Source code in toolbox/io/cp2k.py 
 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
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
def write(self, output_dir=".", fp_params=None, save_dict=False):
    """
    Generate coord.xyz and input.inp for CP2K calculation at output_dir.

    Parameters
    ----------
    output_dir : str
        directory to store coord.xyz and input.inp
    fp_params : dict
        dict for updated parameters
    """
    if fp_params is None:
        fp_params = {}
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    cell = self.atoms.get_cell()
    cell_a = np.array2string(
        cell[0], formatter={"float_kind": lambda x: f"{x:.4f}"}
    )
    cell_a = cell_a[1:-1]
    cell_b = np.array2string(
        cell[1], formatter={"float_kind": lambda x: f"{x:.4f}"}
    )
    cell_b = cell_b[1:-1]
    cell_c = np.array2string(
        cell[2], formatter={"float_kind": lambda x: f"{x:.4f}"}
    )
    cell_c = cell_c[1:-1]

    user_config = fp_params
    update_dict(self.input_dict, user_config)

    if self.input_dict["FORCE_EVAL"].get("QMMM", None) is not None:
        cell_config = {
            "FORCE_EVAL": {
                "SUBSYS": {"CELL": {"A": cell_a, "B": cell_b, "C": cell_c}},
                "QMMM": {
                    "CELL": {
                        "A": cell_a,
                        "B": cell_b,
                        "C": cell_c,
                        "PERIODIC": "XYZ",
                    }
                },
            }
        }
    else:
        cell_config = {
            "FORCE_EVAL": {
                "SUBSYS": {"CELL": {"A": cell_a, "B": cell_b, "C": cell_c}}
            }
        }

    update_dict(self.input_dict, cell_config)
    # output list
    input_str = iterdict(self.input_dict, out_list=["\n"], loop_idx=0)
    # del input_str[0]
    # del input_str[-1]
    # print(input_str)
    str = "\n".join(input_str)
    str = str.strip("\n")

    io.write(os.path.join(output_dir, "coord.xyz"), self.atoms)
    with open(os.path.join(output_dir, "input.inp"), "w", encoding="utf-8") as f:
        f.write(str)

    if save_dict:
        save_dict_json(self.input_dict, os.path.join(output_dir, "input.json"))

Cp2kOutput

Class for parsing CP2K output files.

This class provides methods to extract various properties from CP2K output files, including energies, forces, charges, and other calculation results.

Parameters:

Name Type Description Default
fname str

Path to CP2K output file, by default "output.out"

 'output.out'
ignore_warning bool

Whether to ignore SCF convergence warnings, by default False

 False
Source code in toolbox/io/cp2k.py 
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
class Cp2kOutput:
    """
    Class for parsing CP2K output files.

    This class provides methods to extract various properties from CP2K output files,
    including energies, forces, charges, and other calculation results.

    Parameters
    ----------
    fname : str, optional
        Path to CP2K output file, by default "output.out"
    ignore_warning : bool, optional
        Whether to ignore SCF convergence warnings, by default False
    """

    def __init__(self, fname="output.out", ignore_warning=False) -> None:
        self.output_file = fname
        with open(fname, encoding="UTF-8") as f:
            self.content = f.readlines()
        self.string = "".join(self.content)

        self.check_scf = not ignore_warning
        if (self.check_scf) and (self.scf_loop == -1):
            raise Warning("SCF run NOT converged")

        self.natoms = len(self.atoms)

    @property
    def worktime(self):
        """Get CP2K calculation time.

        Returns
        -------
        float
            Total calculation time in seconds
        """
        # pattern = r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}"
        # out = re.findall(pattern, self.string)
        # t = re.split(r"\.|:|-|\s", out[0])
        # start_time = datetime.datetime(
        #     int(t[0]), int(t[1]), int(t[2]), int(t[3]), int(t[4]), int(t[5]), int(t[6])
        # )
        # t = re.split(r"\.|:|-|\s", out[-1])
        # end_time = datetime.datetime(
        #     int(t[0]), int(t[1]), int(t[2]), int(t[3]), int(t[4]), int(t[5]), int(t[6])
        # )
        # delta_t = end_time - start_time
        # run_time = delta_t.total_seconds()
        # return run_time
        return self.timing_dict["CP2K"]

    def grep_text_search(self, pattern):
        """Search for a pattern in CP2K output file.

        Parameters
        ----------
        pattern : str
            Regular expression pattern to search for

        Returns
        -------
        str
            Line containing the pattern, or empty string if not found
        """
        search_pattern = re.compile(pattern)
        scf_pattern = re.compile(r"SCF run converged in")

        flag = False
        scf_flag = not self.check_scf
        line = None
        for line in self.content:
            line = line.strip("\n")
            if scf_pattern.search(line) is not None:
                scf_flag = True
            if not scf_flag:
                continue
            if search_pattern.search(line) is not None:
                flag = True
                break
        if flag:
            return line
        else:
            return ""

    def grep_texts(self, start_pattern, end_pattern):
        """Extract text between start and end patterns.

        Parameters
        ----------
        start_pattern : str
            Regular expression pattern for start of text block
        end_pattern : str
            Regular expression pattern for end of text block

        Returns
        -------
        tuple
            Tuple containing (nframe, data_lines) where nframe is the number of
            matching blocks and data_lines is a 2D array of the extracted text
        """
        start_pattern = re.compile(start_pattern)
        end_pattern = re.compile(end_pattern)
        scf_pattern = re.compile(r"SCF run converged in")

        flag = False
        scf_flag = not self.check_scf
        data_lines = []
        nframe = 0
        for line in self.content:
            line = line.strip("\n")
            if scf_pattern.search(line) is not None:
                scf_flag = True
            if not scf_flag:
                continue
            if start_pattern.match(line):
                flag = True
            if end_pattern.match(line):
                assert flag is True, (flag, "No data is found in this file.")
                flag = False
                nframe += 1
            if flag is True:
                data_lines.append(line)
        return nframe, np.reshape(data_lines, (nframe, -1))

    def grep_texts_by_nlines(self, start_pattern, nlines):
        """Extract a fixed number of lines after a pattern.

        Parameters
        ----------
        start_pattern : str
            Regular expression pattern to search for
        nlines : int
            Number of lines to extract after the pattern

        Returns
        -------
        tuple
            Tuple containing (nframe, data_lines) where nframe is the number of
            matching blocks and data_lines is a 2D array of the extracted text
        """
        start_pattern = re.compile(start_pattern)
        scf_pattern = re.compile(r"SCF run converged in")

        data_lines = []
        nframe = 0
        scf_flag = not self.check_scf
        for ii, line in enumerate(self.content):
            line = line.strip("\n")
            if scf_pattern.search(line) is not None:
                scf_flag = True
            if not scf_flag:
                continue
            if start_pattern.search(line) is not None:
                data_lines.append(self.content[ii : ii + nlines])
                nframe += 1
                continue
        if nframe == 0:
            raise AttributeError("No data is found in this file.")
        return nframe, np.reshape(data_lines, (nframe, -1))

    @property
    def coord(self):
        """
        Get atomic coordinate from cp2k output.

        Return:
            coord numpy array (n_atom, 3)
        """
        out = parse_init_atomic_coordinates(self.string)
        self.chemical_symbols = out[2]
        return out[0]

    @property
    def atoms(self):
        """Get ASE Atoms object from CP2K output.

        Returns
        -------
        ase.Atoms
            ASE Atoms object with positions, cell, and periodic boundary conditions
        """
        positions = self.coord
        atoms = Atoms(symbols=self.chemical_symbols, positions=positions)
        check_scf = self.check_scf
        self.check_scf = False
        out = parse_all_cells(self.string)
        self.check_scf = check_scf

        atoms.set_cell(out[0])
        atoms.set_pbc(True)
        return atoms

    @property
    def force(self):
        """
        Get atomic force from cp2k output.

        Return:
            force numpy array (n_atom, 3)
        """
        out = parse_atomic_forces_list(self.string)
        return out[0] * AU_TO_EV_EVERY_ANG

    @property
    def energy(self):
        """Get total energy from CP2K output.

        Returns
        -------
        float
            Total energy in eV
        """
        out = parse_energies_list(self.string)
        return out[0] * AU_TO_EV

    @property
    def scf_loop(self):
        """Get number of SCF iterations.

        Returns
        -------
        int
            Number of SCF iterations, or -1 if SCF did not converge
        """
        pattern = r"\s+SCF\srun\sconverged\sin\s+\d+"
        out = re.findall(pattern, self.string)
        if len(out) == 0:
            return -1
        else:
            return int(out[0].split(" ")[-1])

    @property
    def fermi(self):
        """Get Fermi energy from CP2K output.

        Returns
        -------
        float
            Fermi energy in eV
        """
        if self.uks:
            fermi = []
            for line in self.content:
                if re.search("Fermi Energy ", line):
                    fermi.append(float(line.split()[-1]))
            assert (
                len(fermi) == 2
            ), "There should be two Fermi energy in UKS calculation"
            return max(fermi)
        else:
            try:
                pattern = r"\s+Fermi\senergy:\s+.\d\.\d+"
                out = re.findall(pattern, self.string)
                return float(out[0].split(" ")[-1]) * AU_TO_EV
            except IndexError:
                for line in self.content:
                    if re.search("Fermi Energy ", line):
                        return float(line.split()[-1])

    @property
    def m_charge(self):
        """Get Mulliken charges from CP2K output.

        Returns
        -------
        numpy.ndarray
            Array of Mulliken charges for each atom
        """
        start_pattern = "Mulliken Population Analysis"
        nframe, data_lines = self.grep_texts_by_nlines(start_pattern, self.natoms + 3)
        data_list = []
        for line in data_lines[-1, 3:]:
            line_list = line.split()
            data_list.append(float(line_list[-1]))
        return np.reshape(data_list, -1)

    @property
    def h_charge(self):
        """Get Hirshfeld charges from CP2K output.

        Returns
        -------
        numpy.ndarray
            Array of Hirshfeld charges for each atom
        """
        start_pattern = "Hirshfeld Charges"
        _nf, data_lines = self.grep_texts_by_nlines(start_pattern, self.natoms + 3)
        data_list = []
        for line in data_lines[-1, 3:]:
            line_list = line.split()
            data_list.append(float(line_list[-1]))
        return np.reshape(data_list, -1)

    @property
    def dipole_moment(self):
        """Get dipole moment from CP2K output.

        Returns
        -------
        numpy.ndarray
            Dipole moment vector [x, y, z] in Debye
        """
        pattern = "Dipole moment"
        _nf, data_lines = self.grep_texts_by_nlines(pattern, 2)

        data_list = []
        for line in data_lines[-1, 1:]:
            line_list = line.split()
            data_list.append(list(map(float, line_list[1:-1:2])))
        return np.reshape(data_list, 3)

    @property
    def surf_dipole_moment(self):
        """Grep surface dipole moment.

        Total dipole moment perpendicular to
        the slab [electrons-Angstroem]:              -1.5878220042.
        """
        pattern = "Total dipole moment perpendicular to"
        _nf, data_lines = self.grep_texts_by_nlines(pattern, 2)

        line = data_lines[-1, 1]
        line_list = line.split()
        return float(line_list[-1])

    @property
    def potdrop(self):
        """Calculate potential drop across the surface.

        Returns
        -------
        float
            Potential drop in V
        """
        cross_area = np.linalg.norm(np.cross(self.atoms.cell[0], self.atoms.cell[1]))
        DeltaV = self.surf_dipole_moment / cross_area / EPSILON
        return DeltaV

    @property
    def energy_dict(self):
        """Get detailed energy breakdown from CP2K output.

        Returns
        -------
        dict
            Dictionary containing various energy components in eV
        """
        energy_dict = {}

        start_pattern = r"  Total charge density g-space grids:"
        end_pattern = r"  Total energy:"
        _nf, data_lines = self.grep_texts(start_pattern, end_pattern)

        for kw in data_lines[-1, 2:-1]:
            kw = kw.split(":")
            k = kw[0].strip(" ")
            v = float(kw[1]) * AU_TO_EV
            energy_dict[k] = v

        energy_dict.pop("Fermi energy", None)

        tot_e = 0.0
        for v in energy_dict.values():
            tot_e += v
        energy_dict["total"] = tot_e

        return energy_dict

    @property
    def multiplicity(self):
        """Get spin multiplicity from CP2K output.

        Returns
        -------
        int
            Spin multiplicity of the system
        """
        for line in self.content:
            if "Multiplicity" in line:
                break
        return int(line.split()[-1])

    @property
    def uks(self):
        """Check if calculation is spin-unrestricted.

        Returns
        -------
        bool
            True if calculation is spin-unrestricted (UKS), False otherwise
        """
        return any(re.search("Spin unrestricted", line) for line in self.content)

    @property
    def charge(self):
        """Get system charge from CP2K output.

        Returns
        -------
        int
            Total charge of the system
        """
        for line in self.content:
            if "Charge" in line:
                break
        return int(line.split()[-1])

    @property
    def timing_dict(self):
        """Get timing information from CP2K output.

        Returns
        -------
        dict
            Dictionary containing timing information for different CP2K modules
        """
        t_dict = {}
        flag = False
        for line in self.content:
            if "SUBROUTINE" in line:
                flag = True
                continue
            if flag:
                if "---" in line:
                    flag = False
                    break
                data = line.split()
                if len(data) == 7:
                    # print(data)
                    t_dict[data[0]] = float(data[-2])
        return t_dict
atoms property

Get ASE Atoms object from CP2K output.

Returns:

Type Description
Atoms

ASE Atoms object with positions, cell, and periodic boundary conditions
charge property

Get system charge from CP2K output.

Returns:

Type Description
int

Total charge of the system
coord property

Get atomic coordinate from cp2k output.

Return: coord numpy array (n_atom, 3)

dipole_moment property

Get dipole moment from CP2K output.

Returns:

Type Description
ndarray

Dipole moment vector [x, y, z] in Debye
energy property

Get total energy from CP2K output.

Returns:

Type Description
float

Total energy in eV
energy_dict property

Get detailed energy breakdown from CP2K output.

Returns:

Type Description
dict

Dictionary containing various energy components in eV
fermi property

Get Fermi energy from CP2K output.

Returns:

Type Description
float

Fermi energy in eV
force property

Get atomic force from cp2k output.

Return: force numpy array (n_atom, 3)

h_charge property

Get Hirshfeld charges from CP2K output.

Returns:

Type Description
ndarray

Array of Hirshfeld charges for each atom
m_charge property

Get Mulliken charges from CP2K output.

Returns:

Type Description
ndarray

Array of Mulliken charges for each atom
multiplicity property

Get spin multiplicity from CP2K output.

Returns:

Type Description
int

Spin multiplicity of the system
potdrop property

Calculate potential drop across the surface.

Returns:

Type Description
float

Potential drop in V
scf_loop property

Get number of SCF iterations.

Returns:

Type Description
int

Number of SCF iterations, or -1 if SCF did not converge
surf_dipole_moment property

Grep surface dipole moment.

Total dipole moment perpendicular to the slab [electrons-Angstroem]: -1.5878220042.

timing_dict property

Get timing information from CP2K output.

Returns:

Type Description
dict

Dictionary containing timing information for different CP2K modules
uks property

Check if calculation is spin-unrestricted.

Returns:

Type Description
bool

True if calculation is spin-unrestricted (UKS), False otherwise
worktime property

Get CP2K calculation time.

Returns:

Type Description
float

Total calculation time in seconds
grep_text_search(pattern)

Search for a pattern in CP2K output file.

Parameters:

Name Type Description Default
pattern str

Regular expression pattern to search for

 required

Returns:

Type Description
str

Line containing the pattern, or empty string if not found
Source code in toolbox/io/cp2k.py 
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
def grep_text_search(self, pattern):
    """Search for a pattern in CP2K output file.

    Parameters
    ----------
    pattern : str
        Regular expression pattern to search for

    Returns
    -------
    str
        Line containing the pattern, or empty string if not found
    """
    search_pattern = re.compile(pattern)
    scf_pattern = re.compile(r"SCF run converged in")

    flag = False
    scf_flag = not self.check_scf
    line = None
    for line in self.content:
        line = line.strip("\n")
        if scf_pattern.search(line) is not None:
            scf_flag = True
        if not scf_flag:
            continue
        if search_pattern.search(line) is not None:
            flag = True
            break
    if flag:
        return line
    else:
        return ""
grep_texts(start_pattern, end_pattern)

Extract text between start and end patterns.

Parameters:

Name Type Description Default
start_pattern str

Regular expression pattern for start of text block

 required
end_pattern str

Regular expression pattern for end of text block

 required

Returns:

Type Description
tuple

Tuple containing (nframe, data_lines) where nframe is the number of matching blocks and data_lines is a 2D array of the extracted text
Source code in toolbox/io/cp2k.py 
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
def grep_texts(self, start_pattern, end_pattern):
    """Extract text between start and end patterns.

    Parameters
    ----------
    start_pattern : str
        Regular expression pattern for start of text block
    end_pattern : str
        Regular expression pattern for end of text block

    Returns
    -------
    tuple
        Tuple containing (nframe, data_lines) where nframe is the number of
        matching blocks and data_lines is a 2D array of the extracted text
    """
    start_pattern = re.compile(start_pattern)
    end_pattern = re.compile(end_pattern)
    scf_pattern = re.compile(r"SCF run converged in")

    flag = False
    scf_flag = not self.check_scf
    data_lines = []
    nframe = 0
    for line in self.content:
        line = line.strip("\n")
        if scf_pattern.search(line) is not None:
            scf_flag = True
        if not scf_flag:
            continue
        if start_pattern.match(line):
            flag = True
        if end_pattern.match(line):
            assert flag is True, (flag, "No data is found in this file.")
            flag = False
            nframe += 1
        if flag is True:
            data_lines.append(line)
    return nframe, np.reshape(data_lines, (nframe, -1))
grep_texts_by_nlines(start_pattern, nlines)

Extract a fixed number of lines after a pattern.

Parameters:

Name Type Description Default
start_pattern str

Regular expression pattern to search for

 required
nlines int

Number of lines to extract after the pattern

 required

Returns:

Type Description
tuple

Tuple containing (nframe, data_lines) where nframe is the number of matching blocks and data_lines is a 2D array of the extracted text
Source code in toolbox/io/cp2k.py 
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
def grep_texts_by_nlines(self, start_pattern, nlines):
    """Extract a fixed number of lines after a pattern.

    Parameters
    ----------
    start_pattern : str
        Regular expression pattern to search for
    nlines : int
        Number of lines to extract after the pattern

    Returns
    -------
    tuple
        Tuple containing (nframe, data_lines) where nframe is the number of
        matching blocks and data_lines is a 2D array of the extracted text
    """
    start_pattern = re.compile(start_pattern)
    scf_pattern = re.compile(r"SCF run converged in")

    data_lines = []
    nframe = 0
    scf_flag = not self.check_scf
    for ii, line in enumerate(self.content):
        line = line.strip("\n")
        if scf_pattern.search(line) is not None:
            scf_flag = True
        if not scf_flag:
            continue
        if start_pattern.search(line) is not None:
            data_lines.append(self.content[ii : ii + nlines])
            nframe += 1
            continue
    if nframe == 0:
        raise AttributeError("No data is found in this file.")
    return nframe, np.reshape(data_lines, (nframe, -1))

Cp2kPDOS

Class for handling CP2K projected density of states (PDOS) files.

This class provides methods to read and analyze PDOS data from CP2K, including calculation of various electronic properties.

Parameters:

Name Type Description Default
file_name str

Path to PDOS file

 required
Source code in toolbox/io/cp2k.py 
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
class Cp2kPDOS:
    """
    Class for handling CP2K projected density of states (PDOS) files.

    This class provides methods to read and analyze PDOS data from CP2K,
    including calculation of various electronic properties.

    Parameters
    ----------
    file_name : str
        Path to PDOS file
    """

    def __init__(self, file_name: str) -> None:
        self.file_name = file_name
        with open(file_name, encoding="UTF-8") as f:
            line = f.readline()
        # grep the word after "kind"
        self.element = line.split()[line.split().index("kind") + 1]
        self.fermi = float(line.split()[-2]) * AU_TO_EV

        self._data = np.loadtxt(file_name)
        self.energies = self._data[:, 1]
        self.occupation = self._data[:, 2]

        self.dos_data = None
        self.pdos_data = None

    def get_dos(self, broadening: float = 0.01, energy_step: float = 0.01):
        """Ref: https://manual.cp2k.org/trunk/CP2K_INPUT/FORCE_EVAL/PROPERTIES/BANDSTRUCTURE/DOS.html."""
        bin_edges = np.arange(
            self.energies[0], self.energies[-1] + energy_step, energy_step
        )
        bins, dos = gaussian_filter(self.energies, bin_edges, broadening)
        self.dos_data = (bins, dos)
        return self.dos_data

    def get_pdos(self, broadening=0.01, energy_step=0.01, dos_type="total"):
        """Ref: https://manual.cp2k.org/trunk/CP2K_INPUT/FORCE_EVAL/PROPERTIES/BANDSTRUCTURE/DOS.html."""
        bin_edges = np.arange(
            self.energies[0], self.energies[-1] + energy_step, energy_step
        )
        bins, pdos = gaussian_filter(
            self.energies, bin_edges, broadening, weight=self.get_raw_pdos(dos_type)
        )
        self.pdos_data = (bins, pdos)
        return self.pdos_data

    def get_raw_pdos(self, dos_type):
        """Get raw PDOS data of specified type.

        Parameters
        ----------
        dos_type : str
            Type of PDOS to retrieve (e.g., 'total', 's', 'p', 'd', 'f')

        Returns
        -------
        numpy.ndarray
            Raw PDOS data

        Raises
        ------
        NameError
            If the specified PDOS type does not exist
        """
        try:
            return getattr(self, f"_get_raw_pdos_{dos_type}")()
        except AttributeError as e:
            raise NameError("PDOS type does not exist!") from e

    def _get_raw_pdos_total(self):
        return self._data[:, 3:].sum(axis=1)

    def _get_raw_pdos_s(self):
        return self._data[:, 3]

    def _get_raw_pdos_p(self):
        return self._data[:, 4:7].sum(axis=1)

    def _get_raw_pdos_d(self):
        return self._data[:, 7:12].sum(axis=1)

    def _get_raw_pdos_f(self):
        return self._data[:, 12:19].sum(axis=1)

    @property
    def homo(self):
        """Get highest occupied molecular orbital (HOMO) energy.

        Returns
        -------
        float
            HOMO energy relative to Fermi level in eV
        """
        homo_idx = np.where(self.occupation == 0)[0][0] - 1
        return self.energies[homo_idx] - self.fermi

    @property
    def lumo(self):
        """Get lowest unoccupied molecular orbital (LUMO) energy.

        Returns
        -------
        float
            LUMO energy relative to Fermi level in eV
        """
        return self.energies[self.occupation == 0][0] - self.fermi

    @property
    def vbm(self):
        """Get valence band maximum (VBM) energy.

        Returns
        -------
        float
            VBM energy relative to Fermi level in eV
        """
        raw_dos = self.get_raw_pdos("total")
        mask = (self.occupation > 1e-5) & (raw_dos > 1e-3)
        try:
            return self.energies[mask].max() - self.fermi
        except ValueError:
            print("Warning: No VBM is found!")
            return self.energies.min() - self.fermi

    @property
    def cbm(self):
        """Get conduction band minimum (CBM) energy.

        Returns
        -------
        float
            CBM energy relative to Fermi level in eV
        """
        raw_dos = self.get_raw_pdos("total")
        mask = (self.occupation < 1e-5) & (raw_dos > 1e-3)
        try:
            return self.energies[mask].min() - self.fermi
        except ValueError:
            print("Warning: No CBM is found!")
            return self.energies.max() - self.fermi
cbm property

Get conduction band minimum (CBM) energy.

Returns:

Type Description
float

CBM energy relative to Fermi level in eV
homo property

Get highest occupied molecular orbital (HOMO) energy.

Returns:

Type Description
float

HOMO energy relative to Fermi level in eV
lumo property

Get lowest unoccupied molecular orbital (LUMO) energy.

Returns:

Type Description
float

LUMO energy relative to Fermi level in eV
vbm property

Get valence band maximum (VBM) energy.

Returns:

Type Description
float

VBM energy relative to Fermi level in eV
get_dos(broadening=0.01, energy_step=0.01)

Ref: https://manual.cp2k.org/trunk/CP2K_INPUT/FORCE_EVAL/PROPERTIES/BANDSTRUCTURE/DOS.html.

Source code in toolbox/io/cp2k.py 
1739
1740
1741
1742
1743
1744
1745
1746
def get_dos(self, broadening: float = 0.01, energy_step: float = 0.01):
    """Ref: https://manual.cp2k.org/trunk/CP2K_INPUT/FORCE_EVAL/PROPERTIES/BANDSTRUCTURE/DOS.html."""
    bin_edges = np.arange(
        self.energies[0], self.energies[-1] + energy_step, energy_step
    )
    bins, dos = gaussian_filter(self.energies, bin_edges, broadening)
    self.dos_data = (bins, dos)
    return self.dos_data
get_pdos(broadening=0.01, energy_step=0.01, dos_type='total')

Ref: https://manual.cp2k.org/trunk/CP2K_INPUT/FORCE_EVAL/PROPERTIES/BANDSTRUCTURE/DOS.html.

Source code in toolbox/io/cp2k.py 
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
def get_pdos(self, broadening=0.01, energy_step=0.01, dos_type="total"):
    """Ref: https://manual.cp2k.org/trunk/CP2K_INPUT/FORCE_EVAL/PROPERTIES/BANDSTRUCTURE/DOS.html."""
    bin_edges = np.arange(
        self.energies[0], self.energies[-1] + energy_step, energy_step
    )
    bins, pdos = gaussian_filter(
        self.energies, bin_edges, broadening, weight=self.get_raw_pdos(dos_type)
    )
    self.pdos_data = (bins, pdos)
    return self.pdos_data
get_raw_pdos(dos_type)

Get raw PDOS data of specified type.

Parameters:

Name Type Description Default
dos_type str

Type of PDOS to retrieve (e.g., 'total', 's', 'p', 'd', 'f')

 required

Returns:

Type Description
ndarray

Raw PDOS data

Raises:

Type Description
NameError

If the specified PDOS type does not exist
Source code in toolbox/io/cp2k.py 
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
def get_raw_pdos(self, dos_type):
    """Get raw PDOS data of specified type.

    Parameters
    ----------
    dos_type : str
        Type of PDOS to retrieve (e.g., 'total', 's', 'p', 'd', 'f')

    Returns
    -------
    numpy.ndarray
        Raw PDOS data

    Raises
    ------
    NameError
        If the specified PDOS type does not exist
    """
    try:
        return getattr(self, f"_get_raw_pdos_{dos_type}")()
    except AttributeError as e:
        raise NameError("PDOS type does not exist!") from e

MultiFrameCp2kOutput

Bases: Cp2kOutput

Class for parsing CP2K output files with multiple frames.

This class extends Cp2kOutput to handle multi-frame calculations such as molecular dynamics or geometry optimizations.

Parameters:

Name Type Description Default
fname str

Path to CP2K output file, by default "output.out"

 'output.out'
ignore_warning bool

Whether to ignore SCF convergence warnings, by default False

 False
Source code in toolbox/io/cp2k.py 
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
class MultiFrameCp2kOutput(Cp2kOutput):
    """
    Class for parsing CP2K output files with multiple frames.

    This class extends Cp2kOutput to handle multi-frame calculations such as
    molecular dynamics or geometry optimizations.

    Parameters
    ----------
    fname : str, optional
        Path to CP2K output file, by default "output.out"
    ignore_warning : bool, optional
        Whether to ignore SCF convergence warnings, by default False
    """

    def __init__(self, fname="output.out", ignore_warning=False) -> None:
        super().__init__(fname, ignore_warning)

    @property
    def atoms(self):
        """Get ASE Atoms object from multi-frame CP2K output.

        Returns
        -------
        ase.Atoms
            ASE Atoms object with positions, cell, and periodic boundary conditions
        """
        positions = self.coord
        atoms = Atoms(symbols=self.chemical_symbols, positions=positions)
        a = float(self.grep_text_search(r"Vector a").split()[-1])
        b = float(self.grep_text_search(r"Vector b").split()[-1])
        c = float(self.grep_text_search(r"Vector c").split()[-1])
        alpha = float(self.grep_text_search(r"Angle | alpha").split()[-1])
        beta = float(self.grep_text_search(r"Angle | beta").split()[-1])
        gamma = float(self.grep_text_search(r"Angle | gamma").split()[-1])
        atoms.set_cell([a, b, c, alpha, beta, gamma])
        atoms.set_pbc(True)
        return atoms

    @property
    def force(self):
        """
        Get atomic force from cp2k output.

        Return:
            force numpy array (n_atom, 3)
        """
        start_pattern = r" ATOMIC FORCES in"
        end_pattern = r" SUM OF ATOMIC FORCES"
        nframe, data_lines = self.grep_texts(start_pattern, end_pattern)
        data_lines = np.reshape(data_lines, (nframe, -1))

        data_list = []
        for line in data_lines[:, 3:].reshape(-1):
            line_list = line.split()
            data_list.append(
                [
                    float(line_list[3]) * AU_TO_EV_EVERY_ANG,
                    float(line_list[4]) * AU_TO_EV_EVERY_ANG,
                    float(line_list[5]) * AU_TO_EV_EVERY_ANG,
                ]
            )
        return np.reshape(data_list, (nframe, -1, 3))

    @property
    def energy(self):
        """Get total energy from multi-frame CP2K output.

        Returns
        -------
        float
            Total energy in eV
        """
        data = self.grep_text_search("Total energy: ")
        # data = self.grep_text_search("Total FORCE_EVAL")
        data = data.replace("\n", " ")
        data = data.split(" ")
        return float(data[-1]) * AU_TO_EV

    @property
    def scf_loop(self):
        """Get number of SCF iterations from multi-frame CP2K output.

        Returns
        -------
        int
            Number of SCF iterations, or -1 if SCF did not converge
        """
        pattern = r"\s+SCF\srun\sconverged\sin\s+\d+"
        out = re.findall(pattern, self.string)
        if len(out) == 0:
            return -1
        else:
            return int(out[0].split(" ")[-1])

    @property
    def fermi(self):
        """Get Fermi energy from multi-frame CP2K output.

        Returns
        -------
        float
            Fermi energy in eV
        """
        pattern = r"\s+Fermi\sEnergy\s\[eV\]\s:\s+.\d\.\d+"
        out = re.findall(pattern, self.string)
        return float(out[0].split(" ")[-1]) * AU_TO_EV

    @property
    def m_charge(self):
        """Get Mulliken charges from multi-frame CP2K output.

        Returns
        -------
        numpy.ndarray
            Array of Mulliken charges for each atom for each frame
        """
        start_pattern = "Mulliken Population Analysis"
        nframe, data_lines = self.grep_texts_by_nlines(start_pattern, self.natoms + 3)
        data_lines = np.reshape(data_lines, (nframe, -1))

        data_list = []
        for line in data_lines[:, 3:].reshape(-1):
            line_list = line.split()
            data_list.append(float(line_list[-1]))
        return np.reshape(data_list, (nframe, -1))

    @property
    def h_charge(self):
        """Get Hirshfeld charges from multi-frame CP2K output.

        Returns
        -------
        numpy.ndarray
            Array of Hirshfeld charges for each atom for each frame
        """
        start_pattern = "Hirshfeld Charges"
        nframe, data_lines = self.grep_texts_by_nlines(start_pattern, self.natoms + 3)
        data_lines = np.reshape(data_lines, (nframe, -1))

        data_list = []
        for line in data_lines[:, 3:].reshape(-1):
            line_list = line.split()
            data_list.append(float(line_list[-1]))
        return np.reshape(data_list, (nframe, -1))

    @property
    def dipole_moment(self):
        """Get dipole moment from multi-frame CP2K output.

        Returns
        -------
        numpy.ndarray
            Array of dipole moment vectors [x, y, z] in Debye for each frame
        """
        pattern = "Dipole moment"
        nframe, data_lines = self.grep_texts_by_nlines(pattern, 2)
        data_lines = np.reshape(data_lines, (nframe, -1))

        data_list = []
        for line in data_lines[:, 1:].reshape(-1):
            line_list = line.split()
            data_list.append(list(map(float, line_list[1:-1:2])))
        return np.reshape(data_list, (nframe, 3))

    @property
    def surf_dipole_moment(self):
        """Grep surface dipole moment.

        Total dipole moment perpendicular to
        the slab [electrons-Angstroem]:              -1.5878220042.
        """
        pattern = "Total dipole moment perpendicular to"
        nframe, data_lines = self.grep_texts_by_nlines(pattern, 2)
        data_lines = np.reshape(data_lines, (nframe, -1))

        data_list = []
        for line in data_lines[:, 1:].reshape(-1):
            line_list = line.split()
            data_list.append(float(line_list[-1]))
        return np.reshape(data_list, (nframe))

    @property
    def potdrop(self):
        """Calculate potential drop across the surface from multi-frame CP2K output.

        Returns
        -------
        numpy.ndarray
            Array of potential drops in V for each frame
        """
        cross_area = np.linalg.norm(np.cross(self.atoms.cell[0], self.atoms.cell[1]))
        DeltaV = self.surf_dipole_moment / cross_area / EPSILON
        return DeltaV

    @property
    def energy_dict(self):
        """Get detailed energy breakdown from multi-frame CP2K output.

        Returns
        -------
        dict
            Dictionary containing various energy components in eV for each frame
        """
        energy_dict = {}

        start_pattern = r"  Total charge density g-space grids:"
        end_pattern = r"  Total energy:"
        nframe, data_lines = self.grep_texts(start_pattern, end_pattern)
        data_lines = np.reshape(data_lines, (nframe, -1))

        tot_e = 0.0
        for kw in data_lines[-1, 2:-1].reshape(-1):
            kw = kw.split(":")
            k = kw[0].strip(" ")
            v = float(kw[1]) * AU_TO_EV
            energy_dict[k] = [v]
            tot_e += v
        # energy_dict["Total energy"].append(tot_e)
        # for kws in data_lines[1:, 2:-1].reshape(-1):
        #     tot_e = 0.
        #     for kw in kws:
        #         kw = kw.split(":")
        #         k = kw[0].strip(' ')
        #         v = float(kw[1]) * AU_TO_EV
        #         energy_dict[k].append(v)
        #         tot_e += v
        #     # energy_dict["Total energy"].append(tot_e)

        energy_dict.pop("Fermi energy", None)

        return energy_dict
atoms property

Get ASE Atoms object from multi-frame CP2K output.

Returns:

Type Description
Atoms

ASE Atoms object with positions, cell, and periodic boundary conditions
dipole_moment property

Get dipole moment from multi-frame CP2K output.

Returns:

Type Description
ndarray

Array of dipole moment vectors [x, y, z] in Debye for each frame
energy property

Get total energy from multi-frame CP2K output.

Returns:

Type Description
float

Total energy in eV
energy_dict property

Get detailed energy breakdown from multi-frame CP2K output.

Returns:

Type Description
dict

Dictionary containing various energy components in eV for each frame
fermi property

Get Fermi energy from multi-frame CP2K output.

Returns:

Type Description
float

Fermi energy in eV
force property

Get atomic force from cp2k output.

Return: force numpy array (n_atom, 3)

h_charge property

Get Hirshfeld charges from multi-frame CP2K output.

Returns:

Type Description
ndarray

Array of Hirshfeld charges for each atom for each frame
m_charge property

Get Mulliken charges from multi-frame CP2K output.

Returns:

Type Description
ndarray

Array of Mulliken charges for each atom for each frame
potdrop property

Calculate potential drop across the surface from multi-frame CP2K output.

Returns:

Type Description
ndarray

Array of potential drops in V for each frame
scf_loop property

Get number of SCF iterations from multi-frame CP2K output.

Returns:

Type Description
int

Number of SCF iterations, or -1 if SCF did not converge
surf_dipole_moment property

Grep surface dipole moment.

Total dipole moment perpendicular to the slab [electrons-Angstroem]: -1.5878220042.

gaussian_convolve(xs, ys, sigma)

Convolve data with Gaussian kernel.

Parameters:

Name Type Description Default
xs array_like

X values

 required
ys array_like

Y values to convolve

 required
sigma float

Standard deviation of Gaussian

 required

Returns:

Type Description
ndarray

Convolved values
Source code in toolbox/io/cp2k.py 
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
def gaussian_convolve(xs, ys, sigma):
    """Convolve data with Gaussian kernel.

    Parameters
    ----------
    xs : array_like
        X values
    ys : array_like
        Y values to convolve
    sigma : float
        Standard deviation of Gaussian

    Returns
    -------
    np.ndarray
        Convolved values
    """
    len(xs) - 1

    output = []
    for x in xs:
        bins = xs - x
        tmp_out = gaussian_kernel(bins, sigma)
        bin_width = bins[1:] - bins[:-1]
        output.append(
            np.sum(bin_width * ((tmp_out * ys)[1:] + (tmp_out * ys)[:-1]) / 2)
        )
    return np.array(output)

gaussian_kernel(bins, sigma)

Generate Gaussian kernel for convolution.

Parameters:

Name Type Description Default
bins array_like

Array of bin centers

 required
sigma float

Standard deviation of Gaussian

 required

Returns:

Type Description
ndarray

Gaussian kernel values

Raises:

Type Description
AttributeError

If sigma is negative
Source code in toolbox/io/cp2k.py 
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
def gaussian_kernel(bins, sigma):
    """Generate Gaussian kernel for convolution.

    Parameters
    ----------
    bins : array_like
        Array of bin centers
    sigma : float
        Standard deviation of Gaussian

    Returns
    -------
    np.ndarray
        Gaussian kernel values

    Raises
    ------
    AttributeError
        If sigma is negative
    """
    if sigma == 0:
        output = np.zeros_like(bins)
        one_id = np.where(bins == 0.0)[0][0]
        output[one_id] = 1
        return output
    elif sigma > 0:
        A = 1 / (sigma * np.sqrt(2 * np.pi))
        output = np.exp(-bins * bins / (2 * sigma**2))
        output *= A
        return output
    else:
        raise AttributeError("Sigma should be non-negative value.")

generate_ids_list(ids)

Generate compact list representation of IDs.

Parameters:

Name Type Description Default
ids array_like

Array of integer IDs

 required

Returns:

Type Description
list

List of IDs with ranges compressed (e.g., [1, 2..5, 7])
Source code in toolbox/io/cp2k.py 
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
def generate_ids_list(ids):
    """Generate compact list representation of IDs.

    Parameters
    ----------
    ids : array_like
        Array of integer IDs

    Returns
    -------
    list
        List of IDs with ranges compressed (e.g., [1, 2..5, 7])
    """
    ids = np.sort(ids)
    diff = np.diff(ids)
    # ids_std = ""
    ids_list = []
    update_flag = True
    for count, ii in enumerate(diff):
        if update_flag:
            start_id = ids[count]
            update_flag = False
        end_id = ids[count]
        # print(start_id, end_id)
        if ii == 1:
            continue
        else:
            if start_id == end_id:
                # ids_std += "%d" % start_id
                ids_list.append(start_id)
            elif end_id > start_id:
                # ids_std += "%d..%d " % (start_id, end_id)
                ids_list.append(f"{start_id:d}..{end_id:d}")
            else:
                raise ValueError("end_id is smaller than start_id")
            update_flag = True

    if update_flag:
        # ids_std += "%d " % ids[-1]
        ids_list.append(ids[-1])
    else:
        # ids_std += "%d..%d " % (start_id, ids[-1])
        ids_list.append(f"{start_id:d}..{ids[-1]:d}")

    # return ids_std.strip()
    return ids_list

read_wannier_spread(fname='wannier_spread.out')

Read wannier spread file generated by cp2k.

Parameters:

Name Type Description Default
fname str

wannier_spread.out file name

 'wannier_spread.out'

Returns:

Name Type Description
wannier_spread numpy array

wannier spread data
Source code in toolbox/io/cp2k.py 
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
def read_wannier_spread(fname: str = "wannier_spread.out"):
    """
    Read wannier spread file generated by cp2k.

    Parameters
    ----------
    fname : str
        wannier_spread.out file name

    Returns
    -------
    wannier_spread : numpy array
        wannier spread data
    """
    # skip 1, 2, and the last lines
    # save the others as array
    with open(fname, encoding="UTF-8") as f:
        lines = f.readlines()[2:-1]

    wannier = []
    for line in lines:
        # line to array
        line = line.strip().split()
        wannier.append([float(line[1]), float(line[2])])
    return np.array(wannier)

Deep Potential JAX

toolbox.io.dp_jax

DeepMD JAX molecular dynamics module.

This module provides classes for running molecular dynamics simulations using DeepMD JAX implementation with trajectory dumping functionality.

Simulation

Bases: Simulation

DeepMD JAX molecular dynamics simulation class.

This class extends the base DeepMD JAX Simulation class to provide trajectory dumping and logging functionality.

Example
setup simulation

sim = Simulation( model_path="model.pkl", # Has to be an 'energy' or 'dplr' model box=box, # Angstroms type_idx=type_idx, # here the index-element map (e.g. 0-Oxygen, 1-Hydrogen) must match the dataset used to train the model mass=[15.9994, 1.0078, 195.08], # Oxygen, Hydrogen routine="NVT", # 'NVE', 'NVT', 'NPT' (Nosé-Hoover) dt=0.5, # femtoseconds initial_position=initial_position, # Angstroms temperature=330, # Kelvin report_interval=1, # Report every 100 steps seed=np.random.randint(1, 1e5), # Random seed )

sim.run( n_steps, [ TrajDump(atoms, "pos_traj.xyz", 10, append=True), TrajDump(atoms, "vel_traj.xyz", 10, vel=True, append=True), ], )

Source code in toolbox/io/dp_jax.py 
 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
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
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
class Simulation(_Simulation):
    """DeepMD JAX molecular dynamics simulation class.

    This class extends the base DeepMD JAX Simulation class to provide
    trajectory dumping and logging functionality.

    Example
    -------
    # setup simulation
    sim = Simulation(
        model_path="model.pkl",  # Has to be an 'energy' or 'dplr' model
        box=box,  # Angstroms
        type_idx=type_idx,  # here the index-element map (e.g. 0-Oxygen, 1-Hydrogen) must match the dataset used to train the model
        mass=[15.9994, 1.0078, 195.08],  # Oxygen, Hydrogen
        routine="NVT",  # 'NVE', 'NVT', 'NPT' (Nosé-Hoover)
        dt=0.5,  # femtoseconds
        initial_position=initial_position,  # Angstroms
        temperature=330,  # Kelvin
        report_interval=1,  # Report every 100 steps
        seed=np.random.randint(1, 1e5),  # Random seed
    )

    sim.run(
        n_steps,
        [
            TrajDump(atoms, "pos_traj.xyz", 10, append=True),
            TrajDump(atoms, "vel_traj.xyz", 10, vel=True, append=True),
        ],
    )

    """

    def __init__(
        self,
        model_path,
        box,
        type_idx,
        mass,
        routine,
        dt,
        initial_position,
        log_file: Optional[str] = "deepmd_jax.stdout",
        **kwargs,
    ):
        """Initialize Simulation.

        Parameters
        ----------
        model_path : str
            Path to trained model file
        box : array_like
            Simulation box dimensions in Angstroms
        type_idx : array_like
            Index-element mapping for atom types
        mass : array_like
            Atomic masses
        routine : str
            Ensemble type ('NVE', 'NVT', 'NPT')
        dt : float
            Timestep in femtoseconds
        initial_position : array_like
            Initial atomic positions in Angstroms
        log_file : Optional[str], optional
            Log file path, by default "deepmd_jax.stdout"
        **kwargs
            Additional keyword arguments
        """
        super().__init__(
            model_path,
            box,
            type_idx,
            mass,
            routine,
            dt,
            initial_position,
            **kwargs,
        )
        self.log_file = log_file
        if log_file is not None:
            # export all stdout to log_file
            self._stdout = sys.stdout
            with open(log_file, "w", encoding="utf-8") as f:
                sys.stdout = f

    def __del__(self):
        """Restore stdout when object is deleted."""
        if self.log_file is not None:
            sys.stdout.close()
            sys.stdout = self._stdout

    def _initialize_run(self, steps):
        """Initialize simulation run.

        This function resets trajectory for each new run and
        initializes run variables. If the simulation has not
        been run before, it includes the initial state.

        Parameters
        ----------
        steps : int
            Total number of steps to run
        """
        print(f"# Running {steps} steps...")
        self._offset = self.step - int(self._is_initial_state)
        self._tic_of_this_run = time.time()
        self._tic_between_report = time.time()
        self._error_code = 0
        self._print_report()
        self._is_initial_state = False

    def run(self, steps, dump_list: List[TrajDump]):
        """Run the simulation for a number of steps."""
        self._initialize_run(steps)
        remaining_steps = steps
        while remaining_steps > 0:
            # run the simulation for a jit-compiled chunk of steps
            next_chunk = min(
                self.report_interval - self.step % self.report_interval,
                self._step_chunk_size,
                remaining_steps,
            )
            states = (
                self._state,
                self._typed_nbrs if self._static_args["use_neighbor_list"] else None,
                self._error_code,
                self._neighbor_update_profile,
            )
            states_new, traj = self._multiple_inner_step_fn(states, next_chunk)
            state_new, typed_nbrs_new, error_code, profile = states_new
            self._error_code |= error_code

            if self._error_code & 16:
                print("# Warning: Nan or Inf encountered in simulation. Terminating.")
                remaining_steps = 0

            # If there is any hard overflow, we have to re-run the chunk
            if not (
                self._error_code == 0 or self._error_code == 4 or self._error_code == 16
            ):
                self._resolve_error_code()
                continue

            # If nothing overflows, update the tracked state and record the trajectory
            self._keep_nbr_or_lattice_up_to_date()
            self._state = state_new
            self._typed_nbrs = typed_nbrs_new
            self._neighbor_update_profile = profile
            if "NPT" in self._routine:
                self._current_box = self._state.box
            pos_traj, vel_traj, box_traj = traj
            self.step += next_chunk
            remaining_steps -= next_chunk

            # Report at preset regular intervals
            if self.step % self.report_interval == 0 or remaining_steps == 0:
                self._print_report()

            for dump in dump_list:
                if self.step % dump.interval == 0 or remaining_steps == 0:
                    cell = np.concatenate([np.array(box_traj[-1]), [90, 90, 90]])
                    dump.write(pos_traj[-1] if not dump.vel else vel_traj[-1], cell)

        self._print_run_profile(steps, time.time() - self._tic_of_this_run)
        self._keep_nbr_or_lattice_up_to_date()
__del__()

Restore stdout when object is deleted.

Source code in toolbox/io/dp_jax.py 
157
158
159
160
161
def __del__(self):
    """Restore stdout when object is deleted."""
    if self.log_file is not None:
        sys.stdout.close()
        sys.stdout = self._stdout
__init__(model_path, box, type_idx, mass, routine, dt, initial_position, log_file='deepmd_jax.stdout', **kwargs)

Initialize Simulation.

Parameters:

Name Type Description Default
model_path str

Path to trained model file

 required
box array_like

Simulation box dimensions in Angstroms

 required
type_idx array_like

Index-element mapping for atom types

 required
mass array_like

Atomic masses

 required
routine str

Ensemble type ('NVE', 'NVT', 'NPT')

 required
dt float

Timestep in femtoseconds

 required
initial_position array_like

Initial atomic positions in Angstroms

 required
log_file Optional[str]

Log file path, by default "deepmd_jax.stdout"

 'deepmd_jax.stdout'
**kwargs

Additional keyword arguments

 {}
Source code in toolbox/io/dp_jax.py 
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
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
def __init__(
    self,
    model_path,
    box,
    type_idx,
    mass,
    routine,
    dt,
    initial_position,
    log_file: Optional[str] = "deepmd_jax.stdout",
    **kwargs,
):
    """Initialize Simulation.

    Parameters
    ----------
    model_path : str
        Path to trained model file
    box : array_like
        Simulation box dimensions in Angstroms
    type_idx : array_like
        Index-element mapping for atom types
    mass : array_like
        Atomic masses
    routine : str
        Ensemble type ('NVE', 'NVT', 'NPT')
    dt : float
        Timestep in femtoseconds
    initial_position : array_like
        Initial atomic positions in Angstroms
    log_file : Optional[str], optional
        Log file path, by default "deepmd_jax.stdout"
    **kwargs
        Additional keyword arguments
    """
    super().__init__(
        model_path,
        box,
        type_idx,
        mass,
        routine,
        dt,
        initial_position,
        **kwargs,
    )
    self.log_file = log_file
    if log_file is not None:
        # export all stdout to log_file
        self._stdout = sys.stdout
        with open(log_file, "w", encoding="utf-8") as f:
            sys.stdout = f
run(steps, dump_list)

Run the simulation for a number of steps.

Source code in toolbox/io/dp_jax.py 
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
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
def run(self, steps, dump_list: List[TrajDump]):
    """Run the simulation for a number of steps."""
    self._initialize_run(steps)
    remaining_steps = steps
    while remaining_steps > 0:
        # run the simulation for a jit-compiled chunk of steps
        next_chunk = min(
            self.report_interval - self.step % self.report_interval,
            self._step_chunk_size,
            remaining_steps,
        )
        states = (
            self._state,
            self._typed_nbrs if self._static_args["use_neighbor_list"] else None,
            self._error_code,
            self._neighbor_update_profile,
        )
        states_new, traj = self._multiple_inner_step_fn(states, next_chunk)
        state_new, typed_nbrs_new, error_code, profile = states_new
        self._error_code |= error_code

        if self._error_code & 16:
            print("# Warning: Nan or Inf encountered in simulation. Terminating.")
            remaining_steps = 0

        # If there is any hard overflow, we have to re-run the chunk
        if not (
            self._error_code == 0 or self._error_code == 4 or self._error_code == 16
        ):
            self._resolve_error_code()
            continue

        # If nothing overflows, update the tracked state and record the trajectory
        self._keep_nbr_or_lattice_up_to_date()
        self._state = state_new
        self._typed_nbrs = typed_nbrs_new
        self._neighbor_update_profile = profile
        if "NPT" in self._routine:
            self._current_box = self._state.box
        pos_traj, vel_traj, box_traj = traj
        self.step += next_chunk
        remaining_steps -= next_chunk

        # Report at preset regular intervals
        if self.step % self.report_interval == 0 or remaining_steps == 0:
            self._print_report()

        for dump in dump_list:
            if self.step % dump.interval == 0 or remaining_steps == 0:
                cell = np.concatenate([np.array(box_traj[-1]), [90, 90, 90]])
                dump.write(pos_traj[-1] if not dump.vel else vel_traj[-1], cell)

    self._print_run_profile(steps, time.time() - self._tic_of_this_run)
    self._keep_nbr_or_lattice_up_to_date()

TrajDump

Class for dumping trajectory data during simulation.

This class handles writing trajectory data at specified intervals during molecular dynamics simulations.

Source code in toolbox/io/dp_jax.py 
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
66
67
68
69
70
class TrajDump:
    """Class for dumping trajectory data during simulation.

    This class handles writing trajectory data at specified intervals
    during molecular dynamics simulations.
    """

    def __init__(
        self,
        atoms: Atoms,
        fname: str,
        interval: int,
        vel: bool = False,
        **kwargs,
    ) -> None:
        """Initialize TrajDump.

        Parameters
        ----------
        atoms : ase.Atoms
            Atoms object for the system
        fname : str
            Output filename
        interval : int
            Dump interval in timesteps
        vel : bool, optional
            Whether to dump velocities, by default False
        **kwargs
            Additional keyword arguments for writing
        """
        self.fname = fname
        self.interval = interval
        self.vel = vel
        self.atoms = atoms

        self.write_settings = kwargs

    def write(self, positions, cell):
        """Write current frame to file.

        Parameters
        ----------
        positions : array_like
            Atomic positions
        cell : array_like
            Unit cell vectors
        """
        self.atoms.set_positions(positions)
        self.atoms.set_cell(cell)
        io.write(
            self.fname,
            self.atoms,
            **self.write_settings,
        )
__init__(atoms, fname, interval, vel=False, **kwargs)

Initialize TrajDump.

Parameters:

Name Type Description Default
atoms Atoms

Atoms object for the system

 required
fname str

Output filename

 required
interval int

Dump interval in timesteps

 required
vel bool

Whether to dump velocities, by default False

 False
**kwargs

Additional keyword arguments for writing

 {}
Source code in toolbox/io/dp_jax.py 
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
def __init__(
    self,
    atoms: Atoms,
    fname: str,
    interval: int,
    vel: bool = False,
    **kwargs,
) -> None:
    """Initialize TrajDump.

    Parameters
    ----------
    atoms : ase.Atoms
        Atoms object for the system
    fname : str
        Output filename
    interval : int
        Dump interval in timesteps
    vel : bool, optional
        Whether to dump velocities, by default False
    **kwargs
        Additional keyword arguments for writing
    """
    self.fname = fname
    self.interval = interval
    self.vel = vel
    self.atoms = atoms

    self.write_settings = kwargs
write(positions, cell)

Write current frame to file.

Parameters:

Name Type Description Default
positions array_like

Atomic positions

 required
cell array_like

Unit cell vectors

 required
Source code in toolbox/io/dp_jax.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def write(self, positions, cell):
    """Write current frame to file.

    Parameters
    ----------
    positions : array_like
        Atomic positions
    cell : array_like
        Unit cell vectors
    """
    self.atoms.set_positions(positions)
    self.atoms.set_cell(cell)
    io.write(
        self.fname,
        self.atoms,
        **self.write_settings,
    )

DPData I/O

toolbox.io.dpdata

DPData utility module.

This module provides utility functions for working with DPData format, including setting energies, forces, and updating atom types.

set_energy_and_forces(atoms, energy, forces)

Set energy and forces to atoms object.

This function sets the energy and forces on an ASE Atoms object to make atoms.get_potential_energy() and atoms.get_forces() return the given energy and forces.

Parameters:

Name Type Description Default
atoms Atoms

The atoms object to modify

 required
energy float

The energy value to set

 required
forces ndarray

The forces array to set

 required
Source code in toolbox/io/dpdata.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def set_energy_and_forces(atoms: Atoms, energy: float, forces: np.ndarray):
    """Set energy and forces to atoms object.

    This function sets the energy and forces on an ASE Atoms object
    to make atoms.get_potential_energy() and atoms.get_forces()
    return the given energy and forces.

    Parameters
    ----------
    atoms : ase.Atoms
        The atoms object to modify
    energy : float
        The energy value to set
    forces : np.ndarray
        The forces array to set
    """
    calc = SinglePointCalculator(atoms)
    atoms.set_calculator(calc)
    atoms.calc.results["energy"] = energy
    atoms.calc.results["forces"] = forces

update_atype(dname, type_map)

Update atom types in multiple directories.

This function updates atom type files in all subdirectories to match a specified type mapping.

Parameters:

Name Type Description Default
dname str

Directory name to search for type files

 required
type_map List[str]

List of element symbols in desired order

 required
Source code in toolbox/io/dpdata.py 
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
def update_atype(dname: str, type_map: list[str]):
    """Update atom types in multiple directories.

    This function updates atom type files in all subdirectories
    to match a specified type mapping.

    Parameters
    ----------
    dname : str
        Directory name to search for type files
    type_map : List[str]
        List of element symbols in desired order
    """
    fnames = glob.glob(os.path.join(dname, "**/type.raw"), recursive=True)
    fnames.sort()
    for fname in fnames:
        dname = os.path.dirname(fname)
        _type_map = np.loadtxt(os.path.join(dname, "type_map.raw"), dtype=str)
        _atype = np.loadtxt(fname, dtype=int)
        symbols = _type_map[_atype]
        # generate new type.raw according to the input type_map
        assert np.all(np.isin(symbols, type_map)), "Invalid type_map"
        new_atype = np.array([type_map.index(s) for s in symbols], dtype=int)
        np.savetxt(fname, new_atype, fmt="%d")
        np.savetxt(os.path.join(dname, "type_map.raw"), type_map, fmt="%s")

LAMMPS I/O

toolbox.io.lammps

LAMMPS I/O module.

This module provides classes and functions for reading and writing LAMMPS data files, dump files, and log files.

DPLRLammpsData

Bases: LammpsData

Class for DPLR LAMMPS data files.

This class extends LammpsData to handle DPLR (Deep Potential Learning and Reproducing) specific data format.

Example

atoms = io.read("coord.xyz") sel_type = ["O"] sys_charge_dict = { "O": 6.0, "H": 1.0, } lmp_data = DPLRLammpsData(atoms, sel_type=sel_type, sys_charge_dict=sys_charge_dict) lmp_data.write("system.data", format="lammps-data", specorder=["O", "H"], atom_style="full")

Source code in toolbox/io/lammps.py 
296
297
298
299
300
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
class DPLRLammpsData(LammpsData):
    """Class for DPLR LAMMPS data files.

    This class extends LammpsData to handle DPLR (Deep Potential
    Learning and Reproducing) specific data format.

    Example
    -------

    atoms = io.read("coord.xyz")
    sel_type = ["O"]
    sys_charge_dict = {
        "O": 6.0,
        "H": 1.0,
    }
    lmp_data = DPLRLammpsData(atoms, sel_type=sel_type, sys_charge_dict=sys_charge_dict)
    lmp_data.write("system.data", format="lammps-data", specorder=["O", "H"], atom_style="full")
    """

    def __init__(self, atoms, sel_type, sys_charge_dict) -> None:
        atoms, center_ids = self._make_extended_atoms(
            atoms.copy(), sel_type, sys_charge_dict
        )
        super().__init__(atoms)
        # set bonds between real atoms and wannier atoms
        nbonds = len(center_ids)
        bonds = np.ones((nbonds, 4), dtype=int)
        np.copyto(bonds[:, 0], np.arange(nbonds) + 1)
        wannier_ids = np.arange(len(atoms) - nbonds, len(atoms)) + 1
        np.copyto(bonds[:, 2], center_ids)
        np.copyto(bonds[:, 3], wannier_ids)
        self.set_bonds(bonds)

    def _make_extended_atoms(self, atoms, sel_type, sys_charge_dict):
        dummy_type_map = ["He", "Ne", "Ar", "Kr", "Xe", "Rn"]
        # extend Wannier atoms
        center_ids = []
        for ii, _atype in enumerate(sel_type):
            sel_ids = np.where(atoms.symbols == _atype)[0]
            sel_atoms = atoms[sel_ids]
            while dummy_type_map[ii] in atoms.get_chemical_symbols():
                dummy_type_map.pop(ii)
            sel_atoms.symbols[:] = dummy_type_map[ii]
            atoms.extend(sel_atoms)
            center_ids.append(sel_ids + 1)
        center_ids = np.concatenate(center_ids)

        charges = np.full(len(atoms), -8.0)
        for k, v in sys_charge_dict.items():
            charges[atoms.symbols == k] = v
        atoms.set_initial_charges(charges)

        self.dummy_type_map = dummy_type_map[: len(sel_type)]
        return atoms, center_ids

    def write(self, out_file="system.data", specorder=None, **kwargs):
        """Write DPLR LAMMPS data file.

        Parameters
        ----------
        out_file : str, optional
            Output filename, by default "system.data"
        specorder : list, optional
            Specification order for atom types
        **kwargs
            Additional keyword arguments for writing
        """
        if specorder is not None:
            specorder.extend(self.dummy_type_map)
        super().write(out_file, specorder=specorder, **kwargs)
write(out_file='system.data', specorder=None, **kwargs)

Write DPLR LAMMPS data file.

Parameters:

Name Type Description Default
out_file str

Output filename, by default "system.data"

 'system.data'
specorder list

Specification order for atom types

 None
**kwargs

Additional keyword arguments for writing

 {}
Source code in toolbox/io/lammps.py 
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
def write(self, out_file="system.data", specorder=None, **kwargs):
    """Write DPLR LAMMPS data file.

    Parameters
    ----------
    out_file : str, optional
        Output filename, by default "system.data"
    specorder : list, optional
        Specification order for atom types
    **kwargs
        Additional keyword arguments for writing
    """
    if specorder is not None:
        specorder.extend(self.dummy_type_map)
    super().write(out_file, specorder=specorder, **kwargs)

DPLRRestartLammpsData

Bases: LammpsData

Class for DPLR restart LAMMPS data files.

This class extends LammpsData to handle DPLR restart files with specific velocity and topology information.

Example

atoms = io.read("after_system.data", format="lammps-data", sort_by_id=True, Z_of_type={1: 8, 2: 1, 3: 2} ) sel_type = ["O"] lmp_data = DPLRRestartLammpsData(atoms, sel_type=sel_type) lmp_data.write("restart_system.data", format="lammps-data", specorder=["O", "H", "He"], atom_style="full")

Source code in toolbox/io/lammps.py 
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
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
class DPLRRestartLammpsData(LammpsData):
    """Class for DPLR restart LAMMPS data files.

    This class extends LammpsData to handle DPLR restart files
    with specific velocity and topology information.

    Example
    -------

    atoms = io.read("after_system.data",
                format="lammps-data",
                sort_by_id=True,
                Z_of_type={1: 8, 2: 1, 3: 2}
               )
    sel_type = ["O"]
    lmp_data = DPLRRestartLammpsData(atoms, sel_type=sel_type)
    lmp_data.write("restart_system.data", format="lammps-data", specorder=["O", "H", "He"], atom_style="full")
    """

    def __init__(self, atoms, sel_type) -> None:
        n_wannier = np.count_nonzero(np.isin(atoms.symbols, sel_type))
        n_real = len(atoms) - n_wannier
        center_ids = []
        count = n_real
        coords = atoms.get_positions()
        for _atype in sel_type:
            sel_ids = np.where(atoms.symbols == _atype)[0]
            np.copyto(coords[count : count + len(sel_ids)], coords[sel_ids])
            center_ids.append(sel_ids + 1)
            count += len(sel_ids)
        center_ids = np.concatenate(center_ids)
        atoms.set_positions(coords)
        super().__init__(atoms)
        # set velocities
        # from ase internal unit to lammps metal unit (A/ps)
        # https://wiki.fysik.dtu.dk/ase/ase/units.html#units
        vs = atoms.get_velocities() * fs * 1e3
        # vs[n_real:] = 0.0
        self.set_velocities(
            np.concatenate((np.arange(1, len(atoms) + 1).reshape(-1, 1), vs), axis=1)
        )
        # set bonds
        nbonds = len(center_ids)
        bonds = np.ones((nbonds, 4), dtype=int)
        np.copyto(bonds[:, 0], np.arange(nbonds) + 1)
        wannier_ids = np.arange(len(atoms) - nbonds, len(atoms)) + 1
        np.copyto(bonds[:, 2], center_ids)
        np.copyto(bonds[:, 3], wannier_ids)
        self.set_bonds(bonds)

LammpsData

Class for writing LAMMPS data files from ASE Atoms objects.

This class provides functionality to convert ASE Atoms objects to LAMMPS data format with various atom styles.

Source code in toolbox/io/lammps.py 
 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
112
113
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
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
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
class LammpsData:
    """Class for writing LAMMPS data files from ASE Atoms objects.

    This class provides functionality to convert ASE Atoms objects
    to LAMMPS data format with various atom styles.
    """

    def __init__(self, atoms) -> None:
        """Initialize LammpsData.

        Parameters
        ----------
        atoms : ase.Atoms
            ASE Atoms object to convert to LAMMPS data format
        """
        self.atoms = atoms
        self._setup()

        self.angles = None
        self.bonds = None
        self.dihedrals = None
        self.velocities = None

        self.atype = None

    def write(self, out_file="system.data", **kwargs):
        """Write LAMMPS data file.

        Parameters
        ----------
        out_file : str, optional
            Output filename, by default "system.data"
        **kwargs
            Additional keyword arguments for writing
        """
        if self.atype is None:
            # if atype is not set by hand, set by specorder
            specorder = kwargs.get("specorder")
            if specorder is not None:
                self.set_atype_from_specorder(specorder)

        n_atype = len(np.unique(self.atype))
        atom_style = kwargs.get("atom_style", "full")

        with open(out_file, "w", encoding="utf-8") as f:
            header = self._make_header(out_file, n_atype)
            f.write(header)
            body = self._make_atoms(atom_style)
            f.write(body)
            if self.bonds is not None:
                f.write("\nBonds\n\n")
                np.savetxt(f, self.bonds, fmt="%d")
            if self.angles is not None:
                f.write("\nAngles\n\n")
                np.savetxt(f, self.angles, fmt="%d")
            if self.dihedrals is not None:
                f.write("\nDihedrals\n\n")
                np.savetxt(f, self.dihedrals, fmt="%d")
            if self.velocities is not None:
                f.write("\nVelocities\n\n")
                np.savetxt(f, self.velocities, fmt=["%d", "%.16f", "%.16f", "%.16f"])

    def _make_header(self, out_file, n_atype):
        """Generate LAMMPS data file header.

        Parameters
        ----------
        out_file : str
            Output filename
        n_atype : int
            Number of atom types

        Returns
        -------
        str
            Header string for LAMMPS data file
        """
        # cell = self.atoms.cell.cellpar()
        nat = len(self.atoms)
        s = f"{out_file} (written by toolbox by Jia-Xin Zhu)\n\n"
        s += f"{nat} atoms\n"
        s += f"{n_atype} atom types\n"
        if self.bonds is not None:
            s += f"{len(self.bonds)} bonds\n"
            s += f"{len(np.unique(self.bonds[:, 1]))} bond types\n"
        if self.angles is not None:
            s += f"{len(self.angles)} angles\n"
            s += f"{len(np.unique(self.angles[:, 1]))} angle types\n"
        if self.dihedrals is not None:
            s += f"{len(self.dihedrals)} dihedrals\n"
            s += f"{len(np.unique(self.dihedrals[:, 1]))} dihedral types\n"
        # s += "%.4f %.4f xlo xhi\n%.4f %.4f ylo yhi\n%.4f %.4f zlo zhi\n\n\n" % (
        #     0.0, cell[0], 0.0, cell[1], 0.0, cell[2])
        prismobj = Prism(self.atoms.get_cell())
        xhi, yhi, zhi, xy, xz, yz = convert(
            prismobj.get_lammps_prism(), "distance", "ASE", "metal"
        )
        s += f"0.0 {xhi:.6f} xlo xhi\n"
        s += f"0.0 {yhi:.6f} ylo yhi\n"
        s += f"0.0 {zhi:.6f} zlo zhi\n"
        if prismobj.is_skewed():
            s += f"{xy:.6f} {xz:.6f} {yz:.6f} xy xz yz\n"
        s += "\n"
        return s

    def _make_atoms(self, atom_style):
        """Generate atoms section based on atom style.

        Parameters
        ----------
        atom_style : str
            LAMMPS atom style (e.g., "full", "atomic")

        Returns
        -------
        str
            Atoms section string for LAMMPS data file

        Notes
        -----
        Supported styles:
        - full: atom_id res_id type q x y z
        - atomic: atom_id type x y z
        """
        return getattr(self, f"_make_atoms_{atom_style}")()

    def _make_atoms_full(self):
        """Generate atoms section for full atom style.

        Returns
        -------
        str
            Atoms section with format: atom_id res_id type q x y z
        """
        s = "Atoms\n\n"
        for atom in self.atoms:
            ii = atom.index
            s += f"{ii + 1} {self.res_id[ii]} {self.atype[ii]} {self.charges[ii]:.16f} {self.positions[ii][0]:.16f} {self.positions[ii][1]:.16f} {self.positions[ii][2]:.16f}\n"
        return s

    def _make_atoms_atomic(self):
        """Generate atoms section for atomic atom style.

        Returns
        -------
        str
            Atoms section with format: atom_id type x y z
        """
        pass

    def set_res_id(self, res_id):
        """Set residue IDs for atoms.

        Parameters
        ----------
        res_id : array_like
            Array of residue IDs
        """
        self.res_id = np.reshape(res_id, (-1))

    def set_atype(self, atype):
        """Set atom types.

        Parameters
        ----------
        atype : array_like
            Array of atom types
        """
        self.atype = np.reshape(atype, (-1))

    def set_atype_from_specorder(self, specorder):
        """Set atom types from specification order.

        Parameters
        ----------
        specorder : list
            List of element symbols in desired order
        """
        atype = []
        for ii in self.atoms.get_chemical_symbols():
            atype.append(specorder.index(ii))
        self.atype = np.array(atype, dtype=np.int32) + 1

    def set_bonds(self, bonds):
        """Set bonds between atoms.

        Parameters
        ----------
        bonds : array_like
            Array of bonds with shape (n_bonds, 4)
        """
        self.bonds = np.reshape(bonds, (-1, 4))

    def set_angles(self, angles):
        """Set angles between atoms.

        Parameters
        ----------
        angles : array_like
            Array of angles with shape (n_angles, 5)
        """
        self.angles = np.reshape(angles, (-1, 5))

    def set_dihedrals(self, dihedrals):
        """Set dihedrals between atoms.

        Parameters
        ----------
        dihedrals : array_like
            Array of dihedrals with shape (n_dihedrals, 6)
        """
        self.dihedrals = np.reshape(dihedrals, (-1, 6))

    def set_charges(self, charges):
        """Set atomic charges.

        Parameters
        ----------
        charges : array_like
            Array of atomic charges
        """
        self.charges = np.reshape(charges, (-1))

    def set_velocities(self, velocities):
        """Set atomic velocities.

        Parameters
        ----------
        velocities : array_like
            Array of velocities with shape (n_atoms, 4)
        """
        self.velocities = np.reshape(velocities, (-1, 4))

    def _setup(self):
        """Set up initial atom data from ASE Atoms object."""
        self.positions = self.atoms.get_positions()
        if len(self.atoms.get_initial_charges()) > 0:
            self.charges = self.atoms.get_initial_charges().reshape(-1)
        else:
            self.charges = np.zeros(len(self.atoms))

        if hasattr(self, "res_id"):
            assert len(self.res_id) == len(self.atoms)
        else:
            self.res_id = np.zeros(len(self.atoms), dtype=np.int32)
__init__(atoms)

Initialize LammpsData.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object to convert to LAMMPS data format

 required
Source code in toolbox/io/lammps.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def __init__(self, atoms) -> None:
    """Initialize LammpsData.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object to convert to LAMMPS data format
    """
    self.atoms = atoms
    self._setup()

    self.angles = None
    self.bonds = None
    self.dihedrals = None
    self.velocities = None

    self.atype = None
set_angles(angles)

Set angles between atoms.

Parameters:

Name Type Description Default
angles array_like

Array of angles with shape (n_angles, 5)

 required
Source code in toolbox/io/lammps.py 
242
243
244
245
246
247
248
249
250
def set_angles(self, angles):
    """Set angles between atoms.

    Parameters
    ----------
    angles : array_like
        Array of angles with shape (n_angles, 5)
    """
    self.angles = np.reshape(angles, (-1, 5))
set_atype(atype)

Set atom types.

Parameters:

Name Type Description Default
atype array_like

Array of atom types

 required
Source code in toolbox/io/lammps.py 
209
210
211
212
213
214
215
216
217
def set_atype(self, atype):
    """Set atom types.

    Parameters
    ----------
    atype : array_like
        Array of atom types
    """
    self.atype = np.reshape(atype, (-1))
set_atype_from_specorder(specorder)

Set atom types from specification order.

Parameters:

Name Type Description Default
specorder list

List of element symbols in desired order

 required
Source code in toolbox/io/lammps.py 
219
220
221
222
223
224
225
226
227
228
229
230
def set_atype_from_specorder(self, specorder):
    """Set atom types from specification order.

    Parameters
    ----------
    specorder : list
        List of element symbols in desired order
    """
    atype = []
    for ii in self.atoms.get_chemical_symbols():
        atype.append(specorder.index(ii))
    self.atype = np.array(atype, dtype=np.int32) + 1
set_bonds(bonds)

Set bonds between atoms.

Parameters:

Name Type Description Default
bonds array_like

Array of bonds with shape (n_bonds, 4)

 required
Source code in toolbox/io/lammps.py 
232
233
234
235
236
237
238
239
240
def set_bonds(self, bonds):
    """Set bonds between atoms.

    Parameters
    ----------
    bonds : array_like
        Array of bonds with shape (n_bonds, 4)
    """
    self.bonds = np.reshape(bonds, (-1, 4))
set_charges(charges)

Set atomic charges.

Parameters:

Name Type Description Default
charges array_like

Array of atomic charges

 required
Source code in toolbox/io/lammps.py 
262
263
264
265
266
267
268
269
270
def set_charges(self, charges):
    """Set atomic charges.

    Parameters
    ----------
    charges : array_like
        Array of atomic charges
    """
    self.charges = np.reshape(charges, (-1))
set_dihedrals(dihedrals)

Set dihedrals between atoms.

Parameters:

Name Type Description Default
dihedrals array_like

Array of dihedrals with shape (n_dihedrals, 6)

 required
Source code in toolbox/io/lammps.py 
252
253
254
255
256
257
258
259
260
def set_dihedrals(self, dihedrals):
    """Set dihedrals between atoms.

    Parameters
    ----------
    dihedrals : array_like
        Array of dihedrals with shape (n_dihedrals, 6)
    """
    self.dihedrals = np.reshape(dihedrals, (-1, 6))
set_res_id(res_id)

Set residue IDs for atoms.

Parameters:

Name Type Description Default
res_id array_like

Array of residue IDs

 required
Source code in toolbox/io/lammps.py 
199
200
201
202
203
204
205
206
207
def set_res_id(self, res_id):
    """Set residue IDs for atoms.

    Parameters
    ----------
    res_id : array_like
        Array of residue IDs
    """
    self.res_id = np.reshape(res_id, (-1))
set_velocities(velocities)

Set atomic velocities.

Parameters:

Name Type Description Default
velocities array_like

Array of velocities with shape (n_atoms, 4)

 required
Source code in toolbox/io/lammps.py 
272
273
274
275
276
277
278
279
280
def set_velocities(self, velocities):
    """Set atomic velocities.

    Parameters
    ----------
    velocities : array_like
        Array of velocities with shape (n_atoms, 4)
    """
    self.velocities = np.reshape(velocities, (-1, 4))
write(out_file='system.data', **kwargs)

Write LAMMPS data file.

Parameters:

Name Type Description Default
out_file str

Output filename, by default "system.data"

 'system.data'
**kwargs

Additional keyword arguments for writing

 {}
Source code in toolbox/io/lammps.py 
 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
def write(self, out_file="system.data", **kwargs):
    """Write LAMMPS data file.

    Parameters
    ----------
    out_file : str, optional
        Output filename, by default "system.data"
    **kwargs
        Additional keyword arguments for writing
    """
    if self.atype is None:
        # if atype is not set by hand, set by specorder
        specorder = kwargs.get("specorder")
        if specorder is not None:
            self.set_atype_from_specorder(specorder)

    n_atype = len(np.unique(self.atype))
    atom_style = kwargs.get("atom_style", "full")

    with open(out_file, "w", encoding="utf-8") as f:
        header = self._make_header(out_file, n_atype)
        f.write(header)
        body = self._make_atoms(atom_style)
        f.write(body)
        if self.bonds is not None:
            f.write("\nBonds\n\n")
            np.savetxt(f, self.bonds, fmt="%d")
        if self.angles is not None:
            f.write("\nAngles\n\n")
            np.savetxt(f, self.angles, fmt="%d")
        if self.dihedrals is not None:
            f.write("\nDihedrals\n\n")
            np.savetxt(f, self.dihedrals, fmt="%d")
        if self.velocities is not None:
            f.write("\nVelocities\n\n")
            np.savetxt(f, self.velocities, fmt=["%d", "%.16f", "%.16f", "%.16f"])

LammpsDump

Class for writing LAMMPS dump files.

This class provides functionality to write LAMMPS dump files from ASE trajectory objects with proper formatting.

Source code in toolbox/io/lammps.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
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
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
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
class LammpsDump:
    """Class for writing LAMMPS dump files.

    This class provides functionality to write LAMMPS dump files
    from ASE trajectory objects with proper formatting.
    """

    def __init__(self, traj, type_map) -> None:
        """Initialize LammpsDump.

        Parameters
        ----------
        traj : ase.Atoms or list
            Trajectory data as ASE Atoms object or list
        type_map : list or dict
            Type mapping for atoms
        """
        self.traj = traj
        self._set_atype(type_map)
        # self.type_map = type_map

    def write(
        self,
        start=0,
        step=1,
        out_file="out.lammpstrj",
        append=False,
    ):
        """Write LAMMPS dump file.

        Parameters
        ----------
        start : int, optional
            Starting timestep, by default 0
        step : int, optional
            Timestep interval, by default 1
        out_file : str, optional
            Output filename, by default "out.lammpstrj"
        append : bool, optional
            Whether to append to existing file, by default False
        """
        if isinstance(self.traj, Atoms):
            self._write_dump(self.traj, start, out_file, append)
        else:
            nframe = len(self.traj)
            _ts = np.arange(start, step * nframe, step)
            for ts, atoms in zip(_ts, self.traj, strict=False):
                self._write_dump(atoms, ts, out_file, append=True)

    def _set_atype(self, type_map):
        if isinstance(type_map, list):
            self.atype_dict = {}
            for ii, atype in enumerate(type_map, start=1):
                self.atype_dict[atype] = {}
                self.atype_dict[atype]["type"] = ii
                self.atype_dict[atype]["element"] = atype
        elif isinstance(type_map, dict):
            self.atype_dict = type_map
        else:
            raise AttributeError("Unknown type of type_map")

    def _write_dump(self, atoms, ts, out_file, append):
        if append:
            with open(out_file, "a", encoding="utf-8") as f:
                header = self.make_header(atoms, ts)
                f.write(header)
                body = self.make_body(atoms, self.atype_dict)
                f.write(body)
        else:
            with open(out_file, "w", encoding="utf-8") as f:
                header = self.make_header(atoms, ts)
                f.write(header)
                body = self.make_body(atoms, self.atype_dict)
                f.write(body)

    @staticmethod
    def make_header(atoms, ts):
        """Generate LAMMPS dump file header.

        Parameters
        ----------
        atoms : ase.Atoms
            Atoms object containing atomic data
        ts : int
            Current timestep

        Returns
        -------
        str
            Header string for LAMMPS dump file
        """
        bc_dict = {True: "pp", False: "ff"}

        cell = atoms.cell.cellpar()
        bc = atoms.get_pbc()
        nat = len(atoms)
        s = f"ITEM: TIMESTEP\n{ts}\n"
        s += "ITEM: NUMBER OF ATOMS\n"
        s += f"{nat}\n"
        s += f"ITEM: BOX BOUNDS {bc_dict[bc[0]]} {bc_dict[bc[1]]} {bc_dict[bc[2]]}\n"
        s += f"{0.0:.4f} {cell[0]:.4f}\n{0.0:.4f} {cell[1]:.4f}\n{0.0:.4f} {cell[2]:.4f}\n"
        if len(atoms.get_initial_charges()) > 0:
            s += "ITEM: ATOMS id type element x y z q\n"
        else:
            s += "ITEM: ATOMS id type element x y z\n"
        return s

    @staticmethod
    def make_body(atoms, atype_dict):
        """Generate LAMMPS dump file body.

        Parameters
        ----------
        atoms : ase.Atoms
            Atoms object containing atomic data
        atype_dict : dict
            Dictionary mapping atom types to type numbers and elements

        Returns
        -------
        str
            Body string for LAMMPS dump file
        """
        if len(atoms.get_initial_charges()) > 0:
            q_flag = True
            charges = atoms.get_initial_charges()
        else:
            q_flag = False
        ps = atoms.get_positions()

        s = ""
        for atom in atoms:
            ii = atom.index
            if q_flag:
                s += f"{ii + 1} {atype_dict[atom.symbol]['type']} {atype_dict[atom.symbol]['element']} {ps[ii][0]:.16f} {ps[ii][1]:.16f} {ps[ii][2]:.16f} {charges[ii]:.16f}\n"
            else:
                s += f"{ii + 1} {atype_dict[atom.symbol]['type']} {atype_dict[atom.symbol]['element']} {ps[ii][0]:.16f} {ps[ii][1]:.16f} {ps[ii][2]:.16f}\n"
        return s
__init__(traj, type_map)

Initialize LammpsDump.

Parameters:

Name Type Description Default
traj Atoms or list

Trajectory data as ASE Atoms object or list

 required
type_map list or dict

Type mapping for atoms

 required
Source code in toolbox/io/lammps.py 
460
461
462
463
464
465
466
467
468
469
470
471
def __init__(self, traj, type_map) -> None:
    """Initialize LammpsDump.

    Parameters
    ----------
    traj : ase.Atoms or list
        Trajectory data as ASE Atoms object or list
    type_map : list or dict
        Type mapping for atoms
    """
    self.traj = traj
    self._set_atype(type_map)
make_body(atoms, atype_dict) staticmethod

Generate LAMMPS dump file body.

Parameters:

Name Type Description Default
atoms Atoms

Atoms object containing atomic data

 required
atype_dict dict

Dictionary mapping atom types to type numbers and elements

 required

Returns:

Type Description
str

Body string for LAMMPS dump file
Source code in toolbox/io/lammps.py 
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
@staticmethod
def make_body(atoms, atype_dict):
    """Generate LAMMPS dump file body.

    Parameters
    ----------
    atoms : ase.Atoms
        Atoms object containing atomic data
    atype_dict : dict
        Dictionary mapping atom types to type numbers and elements

    Returns
    -------
    str
        Body string for LAMMPS dump file
    """
    if len(atoms.get_initial_charges()) > 0:
        q_flag = True
        charges = atoms.get_initial_charges()
    else:
        q_flag = False
    ps = atoms.get_positions()

    s = ""
    for atom in atoms:
        ii = atom.index
        if q_flag:
            s += f"{ii + 1} {atype_dict[atom.symbol]['type']} {atype_dict[atom.symbol]['element']} {ps[ii][0]:.16f} {ps[ii][1]:.16f} {ps[ii][2]:.16f} {charges[ii]:.16f}\n"
        else:
            s += f"{ii + 1} {atype_dict[atom.symbol]['type']} {atype_dict[atom.symbol]['element']} {ps[ii][0]:.16f} {ps[ii][1]:.16f} {ps[ii][2]:.16f}\n"
    return s
make_header(atoms, ts) staticmethod

Generate LAMMPS dump file header.

Parameters:

Name Type Description Default
atoms Atoms

Atoms object containing atomic data

 required
ts int

Current timestep

 required

Returns:

Type Description
str

Header string for LAMMPS dump file
Source code in toolbox/io/lammps.py 
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
@staticmethod
def make_header(atoms, ts):
    """Generate LAMMPS dump file header.

    Parameters
    ----------
    atoms : ase.Atoms
        Atoms object containing atomic data
    ts : int
        Current timestep

    Returns
    -------
    str
        Header string for LAMMPS dump file
    """
    bc_dict = {True: "pp", False: "ff"}

    cell = atoms.cell.cellpar()
    bc = atoms.get_pbc()
    nat = len(atoms)
    s = f"ITEM: TIMESTEP\n{ts}\n"
    s += "ITEM: NUMBER OF ATOMS\n"
    s += f"{nat}\n"
    s += f"ITEM: BOX BOUNDS {bc_dict[bc[0]]} {bc_dict[bc[1]]} {bc_dict[bc[2]]}\n"
    s += f"{0.0:.4f} {cell[0]:.4f}\n{0.0:.4f} {cell[1]:.4f}\n{0.0:.4f} {cell[2]:.4f}\n"
    if len(atoms.get_initial_charges()) > 0:
        s += "ITEM: ATOMS id type element x y z q\n"
    else:
        s += "ITEM: ATOMS id type element x y z\n"
    return s
write(start=0, step=1, out_file='out.lammpstrj', append=False)

Write LAMMPS dump file.

Parameters:

Name Type Description Default
start int

Starting timestep, by default 0

 0
step int

Timestep interval, by default 1

 1
out_file str

Output filename, by default "out.lammpstrj"

 'out.lammpstrj'
append bool

Whether to append to existing file, by default False

 False
Source code in toolbox/io/lammps.py 
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
500
def write(
    self,
    start=0,
    step=1,
    out_file="out.lammpstrj",
    append=False,
):
    """Write LAMMPS dump file.

    Parameters
    ----------
    start : int, optional
        Starting timestep, by default 0
    step : int, optional
        Timestep interval, by default 1
    out_file : str, optional
        Output filename, by default "out.lammpstrj"
    append : bool, optional
        Whether to append to existing file, by default False
    """
    if isinstance(self.traj, Atoms):
        self._write_dump(self.traj, start, out_file, append)
    else:
        nframe = len(self.traj)
        _ts = np.arange(start, step * nframe, step)
        for ts, atoms in zip(_ts, self.traj, strict=False):
            self._write_dump(atoms, ts, out_file, append=True)

LammpsLog

Class for parsing LAMMPS log files.

This class provides functionality to extract information from LAMMPS log files, including timing and performance data.

Source code in toolbox/io/lammps.py 
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
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
class LammpsLog:
    """Class for parsing LAMMPS log files.

    This class provides functionality to extract information
    from LAMMPS log files, including timing and performance data.
    """

    def __init__(self, fname="log.lammps") -> None:
        """Initialize LammpsLog.

        Parameters
        ----------
        fname : str, optional
            Path to LAMMPS log file, by default "log.lammps"
        """
        self.log_file = fname
        with open(fname) as f:
            self.content = f.readlines()
        # self.string = "".join(self.content)
        self.setup()

    def setup(self):
        """Parse log file to extract performance information."""
        for line in self.content:
            if re.search("MPI tasks", line):
                self.cpu_util = float(line.split()[0][:-1])
                self.n_mpi = int(line.split()[4])
                self.n_thread = int(line.split()[-3])
            if re.match("Loop time of", line):
                self.n_atoms = int(line.split()[-2])
                self.n_step = int(line.split()[-5])
                self.n_proc = int(line.split()[5])
                self.wall_time = float(line.split()[3])
            if re.match("Performance", line):
                out = line.split()
                self.performance = {
                    "ns/day": float(out[1]),
                    "h/ns": float(out[3]),
                    "timesteps/s": float(out[5]),
                }
                try:
                    self.performance["katom-step/s"] = float(out[7])
                except IndexError:
                    self.performance["katom-step/s"] = (
                        float(out[5]) * self.n_atoms / 1e3
                    )

    @property
    def timing_breakdown(self):
        """Get timing breakdown from LAMMPS log file.

        Returns
        -------
        dict
            Dictionary containing timing information for different LAMMPS operations
            with time spent and percentage of total time
        """
        start = False
        timing_breakdown = []
        for line in self.content:
            if start:
                timing_breakdown.append(line)
            if re.match("MPI task timing breakdown", line):
                start = True
            if start and re.match("Nlocal:", line):
                break

        ldata = [line.split("|") for line in timing_breakdown[2:-2]]
        # read data into a dict
        data = {
            d[0]
            .strip()
            .lower(): {
                "time": float(d[2].strip()),
                "percentage": float(d[-1].strip()),
            }
            for d in ldata
        }
        return data
timing_breakdown property

Get timing breakdown from LAMMPS log file.

Returns:

Type Description
dict

Dictionary containing timing information for different LAMMPS operations with time spent and percentage of total time
__init__(fname='log.lammps')

Initialize LammpsLog.

Parameters:

Name Type Description Default
fname str

Path to LAMMPS log file, by default "log.lammps"

 'log.lammps'
Source code in toolbox/io/lammps.py 
600
601
602
603
604
605
606
607
608
609
610
611
612
def __init__(self, fname="log.lammps") -> None:
    """Initialize LammpsLog.

    Parameters
    ----------
    fname : str, optional
        Path to LAMMPS log file, by default "log.lammps"
    """
    self.log_file = fname
    with open(fname) as f:
        self.content = f.readlines()
    # self.string = "".join(self.content)
    self.setup()
setup()

Parse log file to extract performance information.

Source code in toolbox/io/lammps.py 
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
def setup(self):
    """Parse log file to extract performance information."""
    for line in self.content:
        if re.search("MPI tasks", line):
            self.cpu_util = float(line.split()[0][:-1])
            self.n_mpi = int(line.split()[4])
            self.n_thread = int(line.split()[-3])
        if re.match("Loop time of", line):
            self.n_atoms = int(line.split()[-2])
            self.n_step = int(line.split()[-5])
            self.n_proc = int(line.split()[5])
            self.wall_time = float(line.split()[3])
        if re.match("Performance", line):
            out = line.split()
            self.performance = {
                "ns/day": float(out[1]),
                "h/ns": float(out[3]),
                "timesteps/s": float(out[5]),
            }
            try:
                self.performance["katom-step/s"] = float(out[7])
            except IndexError:
                self.performance["katom-step/s"] = (
                    float(out[5]) * self.n_atoms / 1e3
                )

OHHWaterLammpsData

Bases: LammpsData

Class for O-H-H water LAMMPS data files.

This class extends LammpsData to specifically handle water molecules with O-H-H topology, automatically generating bonds and angles.

Source code in toolbox/io/lammps.py 
419
420
421
422
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
class OHHWaterLammpsData(LammpsData):
    """Class for O-H-H water LAMMPS data files.

    This class extends LammpsData to specifically handle water molecules
    with O-H-H topology, automatically generating bonds and angles.
    """

    def __init__(self, atoms) -> None:
        """Initialize OHHWaterLammpsData.

        Parameters
        ----------
        atoms : ase.Atoms
            ASE Atoms object containing water molecules
        """
        super().__init__(atoms)

        oxygen_ids = np.where(atoms.symbols == "O")[0] + 1
        hydrogen_ids = np.where(atoms.symbols == "H")[0] + 1

        bonds, angles = generate_water_bonds_and_angles(
            oxygen_ids,
            hydrogen_ids,
        )

        n_water = len(oxygen_ids)
        res_id = np.arange(n_water) + 1
        res_id = np.tile(res_id.reshape(-1, 1), [1, 3]).reshape(-1)

        self.set_bonds(bonds)
        self.set_angles(angles)
        self.set_res_id(res_id)
__init__(atoms)

Initialize OHHWaterLammpsData.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object containing water molecules

 required
Source code in toolbox/io/lammps.py 
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
def __init__(self, atoms) -> None:
    """Initialize OHHWaterLammpsData.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object containing water molecules
    """
    super().__init__(atoms)

    oxygen_ids = np.where(atoms.symbols == "O")[0] + 1
    hydrogen_ids = np.where(atoms.symbols == "H")[0] + 1

    bonds, angles = generate_water_bonds_and_angles(
        oxygen_ids,
        hydrogen_ids,
    )

    n_water = len(oxygen_ids)
    res_id = np.arange(n_water) + 1
    res_id = np.tile(res_id.reshape(-1, 1), [1, 3]).reshape(-1)

    self.set_bonds(bonds)
    self.set_angles(angles)
    self.set_res_id(res_id)

generate_water_bonds_and_angles(oxygen_ids, hydrogen_ids)

Generate bonds and angles for water molecules.

This function creates bonds and angles for water molecules based on O-H-H topology.

Parameters:

Name Type Description Default
oxygen_ids array_like

Array of oxygen atom indices (1-based)

 required
hydrogen_ids array_like

Array of hydrogen atom indices (1-based)

 required

Returns:

Type Description
tuple

Tuple of (bonds, angles) where: - bonds: array with shape (n_water*2, 4) - angles: array with shape (n_water, 5)
Source code in toolbox/io/lammps.py 
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
def generate_water_bonds_and_angles(oxygen_ids, hydrogen_ids):
    """Generate bonds and angles for water molecules.

    This function creates bonds and angles for water molecules
    based on O-H-H topology.

    Parameters
    ----------
    oxygen_ids : array_like
        Array of oxygen atom indices (1-based)
    hydrogen_ids : array_like
        Array of hydrogen atom indices (1-based)

    Returns
    -------
    tuple
        Tuple of (bonds, angles) where:
        - bonds: array with shape (n_water*2, 4)
        - angles: array with shape (n_water, 5)
    """
    oxygen_ids = np.array(oxygen_ids)
    hydrogen_ids = np.array(hydrogen_ids)
    n_water = len(oxygen_ids)

    nbonds = n_water * 2
    bonds = np.ones((nbonds, 4), dtype=int)
    np.copyto(bonds[:, 0], np.arange(nbonds) + 1)
    np.copyto(bonds[::2, 2], oxygen_ids)
    np.copyto(bonds[1::2, 2], oxygen_ids)
    np.copyto(bonds[:, 3], hydrogen_ids)

    angles = np.ones((n_water, 5), dtype=int)
    np.copyto(angles[:, 0], np.arange(n_water) + 1)
    np.copyto(angles[:, 2], hydrogen_ids[::2])
    np.copyto(angles[:, 3], oxygen_ids)
    np.copyto(angles[:, 4], hydrogen_ids[1::2])
    return bonds, angles

make_dump_body(atoms, atype_dict)

Generate LAMMPS dump file body.

Parameters:

Name Type Description Default
atoms Atoms

Atoms object

 required
atype_dict dict

Atom type dictionary

 required

Returns:

Type Description
str

Body string for LAMMPS dump file
Source code in toolbox/io/lammps.py 
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
def make_dump_body(atoms, atype_dict):
    """Generate LAMMPS dump file body.

    Parameters
    ----------
    atoms : ase.Atoms
        Atoms object
    atype_dict : dict
        Atom type dictionary

    Returns
    -------
    str
        Body string for LAMMPS dump file
    """
    if len(atoms.get_initial_charges()) > 0:
        q_flag = True
        charges = atoms.get_initial_charges()
    else:
        q_flag = False
    ps = atoms.get_positions()

    s = ""
    for atom in atoms:
        ii = atom.index
        if q_flag:
            s += "{:d} {:d} {:s} {:.16f} {:.16f} {:.16f} {:.16f}\n".format(
                ii + 1,
                atype_dict[atom.symbol]["type"],
                atype_dict[atom.symbol]["element"],
                ps[ii][0],
                ps[ii][1],
                ps[ii][2],
                charges[ii],
            )
        else:
            s += "{:d} {:d} {:s} {:.16f} {:.16f} {:.16f}\n".format(
                ii + 1,
                atype_dict[atom.symbol]["type"],
                atype_dict[atom.symbol]["element"],
                ps[ii][0],
                ps[ii][1],
                ps[ii][2],
            )
    return s

make_dump_header(atoms, ts)

Generate LAMMPS dump file header.

Parameters:

Name Type Description Default
atoms Atoms

Atoms object

 required
ts int

Current timestep

 required

Returns:

Type Description
str

Header string for LAMMPS dump file
Source code in toolbox/io/lammps.py 
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
def make_dump_header(atoms, ts):
    """Generate LAMMPS dump file header.

    Parameters
    ----------
    atoms : ase.Atoms
        Atoms object
    ts : int
        Current timestep

    Returns
    -------
    str
        Header string for LAMMPS dump file
    """
    cell = atoms.cell.cellpar()
    nat = len(atoms)
    s = f"ITEM: TIMESTEP\n{ts:d}\n"
    s += "ITEM: NUMBER OF ATOMS\n"
    s += f"{nat:d}\n"
    s += "ITEM: BOX BOUNDS pp pp pp\n"
    s += f"{0.0:.4f} {cell[0]:.4f}\n{0.0:.4f} {cell[1]:.4f}\n{0.0:.4f} {cell[2]:.4f}\n"
    if len(atoms.get_initial_charges()) > 0:
        s += "ITEM: ATOMS id type element x y z q\n"
    else:
        s += "ITEM: ATOMS id type element x y z\n"
    return s

read_dump(dump_file='dump.lammpstrj')

Read LAMMPS dump file and extract trajectory data.

Parameters:

Name Type Description Default
dump_file str

Path to LAMMPS dump file, by default "dump.lammpstrj"

 'dump.lammpstrj'

Returns:

Type Description
tuple

Tuple of (coords, forces, boxs, type_list) where: - coords: array of atomic positions with shape (n_frames, n_atoms*3) - forces: array of atomic forces with shape (n_frames, n_atoms*3) - boxs: array of box vectors with shape (n_frames, 9) - type_list: array of atom types with shape (n_atoms,)
Source code in toolbox/io/lammps.py 
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
def read_dump(dump_file="dump.lammpstrj"):
    """Read LAMMPS dump file and extract trajectory data.

    Parameters
    ----------
    dump_file : str, optional
        Path to LAMMPS dump file, by default "dump.lammpstrj"

    Returns
    -------
    tuple
        Tuple of (coords, forces, boxs, type_list) where:
        - coords: array of atomic positions with shape (n_frames, n_atoms*3)
        - forces: array of atomic forces with shape (n_frames, n_atoms*3)
        - boxs: array of box vectors with shape (n_frames, 9)
        - type_list: array of atom types with shape (n_atoms,)
    """
    traj = io.read(dump_file, index=":")
    coords = []
    forces = []
    boxs = []
    for atoms in traj:
        forces.append(atoms.get_forces())
        coords.append(atoms.get_positions())
        boxs.append(atoms.get_cell())
    coords = np.reshape(coords, (len(traj), -1))
    forces = np.reshape(forces, (len(traj), -1))
    boxs = np.reshape(boxs, (len(traj), -1))
    type_list = atoms.get_array("numbers")
    type_list = np.array(type_list) - 1
    return coords, forces, boxs, type_list

write_dump(traj, type_map, start=0, step=1, out_file='out.lammpstrj', append=False)

Write LAMMPS dump file from trajectory.

Parameters:

Name Type Description Default
traj Atoms or list

Trajectory data

 required
type_map list or dict

Type mapping for atoms

 required
start int

Starting timestep, by default 0

 0
step int

Timestep interval, by default 1

 1
out_file str

Output filename, by default "out.lammpstrj"

 'out.lammpstrj'
append bool

Whether to append to existing file, by default False

 False
Source code in toolbox/io/lammps.py 
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
def write_dump(traj, type_map, start=0, step=1, out_file="out.lammpstrj", append=False):
    """Write LAMMPS dump file from trajectory.

    Parameters
    ----------
    traj : ase.Atoms or list
        Trajectory data
    type_map : list or dict
        Type mapping for atoms
    start : int, optional
        Starting timestep, by default 0
    step : int, optional
        Timestep interval, by default 1
    out_file : str, optional
        Output filename, by default "out.lammpstrj"
    append : bool, optional
        Whether to append to existing file, by default False
    """
    if isinstance(type_map, list):
        atype_dict = {}
        for ii, atype in enumerate(type_map, start=1):
            atype_dict[atype] = {}
            atype_dict[atype]["type"] = ii
            atype_dict[atype]["element"] = atype
    elif isinstance(type_map, dict):
        atype_dict = type_map
    else:
        raise AttributeError("Unknown type of type_map")

    if isinstance(traj, Atoms):
        _write_dump(traj, atype_dict, start, out_file, append)
    else:
        nframe = len(traj)
        _ts = np.arange(start, step * nframe, step)
        for ts, atoms in zip(_ts, traj, strict=False):
            _write_dump(atoms, atype_dict, ts, out_file, append=True)

Plotting Module

Core Plotting Functions

toolbox.plot.core

Core plotting utilities module.

This module provides core plotting functions and utilities for creating various types of plots including learning curves, RMSE plots, binned statistics, and colored line plots.

ax_bin_stats(ax, x, y, bins)

Create binned statistics plot on given axes.

Parameters:

Name Type Description Default
ax Axes

Matplotlib axes object

 required
x array_like

X values

 required
y array_like

Y values

 required
bins int

Number of bins

 required

Returns:

Type Description
dict

Dictionary containing binned statistics data
Source code in toolbox/plot/core.py 
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
220
221
def ax_bin_stats(ax, x, y, bins):
    """Create binned statistics plot on given axes.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Matplotlib axes object
    x : array_like
        X values
    y : array_like
        Y values
    bins : int
        Number of bins

    Returns
    -------
    dict
        Dictionary containing binned statistics data
    """
    # ref line (y = 0)
    ax.axhline(y=0, color="gray")
    # mean
    bin_means, bin_edges, binnumber = stats.binned_statistic(x, y, bins=bins)
    bin_means[np.isnan(bin_means)] = 0.0
    bin_centers = bin_edges[1:] - (bin_edges[1] - bin_edges[0]) / 2
    ax.plot(bin_centers, bin_means, color="black", lw=1.5, label="mean")

    # max/min
    bin_maxs, bin_edges, binnumber = stats.binned_statistic(
        x, y, statistic="max", bins=bins
    )
    bin_mins, bin_edges, binnumber = stats.binned_statistic(
        x, y, statistic="min", bins=bins
    )
    bin_maxs[np.isnan(bin_maxs)] = 0.0
    bin_mins[np.isnan(bin_maxs)] = 0.0
    ax.fill_between(bin_centers, bin_mins, bin_maxs, color="silver", label="[min, max]")
    # std
    bin_stds, bin_edges, binnumber = stats.binned_statistic(
        x, y, statistic="std", bins=bins
    )
    bin_stds[np.isnan(bin_stds)] = 0.0
    ax.fill_between(
        bin_centers,
        bin_means - bin_stds,
        bin_means + bin_stds,
        color="gray",
        label="[mean-std, mean+std]",
    )
    data = {
        "grid": bin_centers,
        "mean": bin_means,
        "max": bin_maxs,
        "min": bin_mins,
        "std": bin_stds,
    }
    return data

ax_colormap_lines(ax, xs, ys, labels, scale=(0.0, 1.0), colormap='GnBu', fmt='%f', **kwargs)

Plot multiple colored lines on axes.

Parameters:

Name Type Description Default
ax Axes

Matplotlib axes object

 required
xs list

List of x data arrays

 required
ys list

List of y data arrays

 required
labels list

List of legend labels

 required
scale tuple

Scale range for normalization, by default (0.0, 1.0)

 (0.0, 1.0)
colormap str

Matplotlib colormap name, by default "GnBu"

 'GnBu'
fmt str

Format string for legend labels, by default "%f"

 '%f'
**kwargs

Additional keyword arguments passed to ax.plot

 {}

Returns:

Type Description
None

Function modifies the axes object in place
Source code in toolbox/plot/core.py 
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
def ax_colormap_lines(
    ax, xs, ys, labels, scale=(0.0, 1.0), colormap="GnBu", fmt="%f", **kwargs
):
    """Plot multiple colored lines on axes.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Matplotlib axes object
    xs : list
        List of x data arrays
    ys : list
        List of y data arrays
    labels : list
        List of legend labels
    scale : tuple, optional
        Scale range for normalization, by default (0.0, 1.0)
    colormap : str, optional
        Matplotlib colormap name, by default "GnBu"
    fmt : str, optional
        Format string for legend labels, by default "%f"
    **kwargs
        Additional keyword arguments passed to ax.plot

    Returns
    -------
    None
        Function modifies the axes object in place
    """
    labels = np.array(labels)
    # normalization
    labels = (labels - labels.min()) / (labels.max() - labels.min())
    cm_scales = (labels - scale[0]) / (scale[1] - scale[0])
    for x, y, label, cm_scale in zip(xs, ys, labels, cm_scales, strict=False):
        ax.plot(
            x, y, color=plt.get_cmap(colormap)(cm_scale), label=fmt % label, **kwargs
        )
    ax.set_xlim(np.min(x), np.max(x))

ax_rmse(ax, x, y)

Create RMSE scatter plot with reference line.

This function creates a scatter plot of data points with a reference line (y=x) for visual comparison.

Parameters:

Name Type Description Default
ax Axes

Matplotlib axes object

 required
x array_like

X values

 required
y array_like

Y values

 required
Source code in toolbox/plot/core.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def ax_rmse(ax, x, y):
    """Create RMSE scatter plot with reference line.

    This function creates a scatter plot of data points
    with a reference line (y=x) for visual comparison.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Matplotlib axes object
    x : array_like
        X values
    y : array_like
        Y values
    """
    # scatter
    ax.scatter(x, y, color="steelblue", alpha=0.2)
    # ref line
    ref = np.arange(x.min(), x.max(), (x.max() - x.min()) / 100)
    ax.plot(ref, ref, color="firebrick", lw=1.5)

ax_setlabel(ax, xlabel, ylabel, **kwargs)

Set axis labels for matplotlib axes.

Parameters:

Name Type Description Default
ax Axes

Matplotlib axes object

 required
xlabel str

X-axis label

 required
ylabel str

Y-axis label

 required
**kwargs

Additional keyword arguments passed to set_xlabel/set_ylabel

 {}
Source code in toolbox/plot/core.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def ax_setlabel(ax, xlabel, ylabel, **kwargs):
    """Set axis labels for matplotlib axes.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Matplotlib axes object
    xlabel : str
        X-axis label
    ylabel : str
        Y-axis label
    **kwargs
        Additional keyword arguments passed to set_xlabel/set_ylabel
    """
    ax.set_xlabel(xlabel, **kwargs)
    ax.set_ylabel(ylabel, **kwargs)

plot_bin_stats(x, y, xlabel, ylabel, bins=None)

Plot binned statistics with confidence intervals.

This function creates a plot showing mean, standard deviation, and min/max ranges for binned data.

Parameters:

Name Type Description Default
x array_like

X values

 required
y array_like

Y values

 required
xlabel str

X-axis label

 required
ylabel str

Y-axis label

 required
bins int

Number of bins, by default len(x)//10

 None

Returns:

Type Description
tuple

Tuple of (fig, ax) containing matplotlib figure and axes
Source code in toolbox/plot/core.py 
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
def plot_bin_stats(x, y, xlabel, ylabel, bins=None):
    """Plot binned statistics with confidence intervals.

    This function creates a plot showing mean, standard deviation,
    and min/max ranges for binned data.

    Parameters
    ----------
    x : array_like
        X values
    y : array_like
        Y values
    xlabel : str
        X-axis label
    ylabel : str
        Y-axis label
    bins : int, optional
        Number of bins, by default len(x)//10

    Returns
    -------
    tuple
        Tuple of (fig, ax) containing matplotlib figure and axes
    """
    x = np.array(x).flatten()
    y = np.array(y).flatten()
    if bins is None:
        bins = len(x) // 10

    fig, ax = plt.subplots(figsize=[6, 4], dpi=200)
    ax_bin_stats(ax, x, y, bins=bins)
    ax_setlabel(ax, xlabel, ylabel)
    ax.legend(loc="center left", bbox_to_anchor=(1.1, 0.5))
    return fig, ax

plot_colormap_lines(xs, ys, legends, xlabel, ylabel, colormap='GnBu')

Plot multiple lines with colormap.

This function creates a plot with multiple lines colored according to a colormap.

Parameters:

Name Type Description Default
xs list

List of x data arrays

 required
ys list

List of y data arrays

 required
legends list

List of legend labels

 required
xlabel str

X-axis label

 required
ylabel str

Y-axis label

 required
colormap str

Matplotlib colormap name, by default "GnBu"

 'GnBu'

Returns:

Type Description
tuple

Tuple of (fig, ax) containing matplotlib figure and axes
Source code in toolbox/plot/core.py 
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
def plot_colormap_lines(xs, ys, legends, xlabel, ylabel, colormap="GnBu"):
    """Plot multiple lines with colormap.

    This function creates a plot with multiple lines colored
    according to a colormap.

    Parameters
    ----------
    xs : list
        List of x data arrays
    ys : list
        List of y data arrays
    legends : list
        List of legend labels
    xlabel : str
        X-axis label
    ylabel : str
        Y-axis label
    colormap : str, optional
        Matplotlib colormap name, by default "GnBu"

    Returns
    -------
    tuple
        Tuple of (fig, ax) containing matplotlib figure and axes
    """
    fig, ax = plt.subplots(figsize=[6, 4], dpi=200)
    ax_colormap_lines(ax, xs, ys, legends, colormap)
    ax_setlabel(ax, xlabel, ylabel)
    ax.legend(loc="center left", bbox_to_anchor=(1.1, 0.5))
    return fig, ax

plot_lcurve(fname, col, xlabel=None, ylabel=None, **kwargs)

Plot learning curve from data file.

This function reads a data file and plots a learning curve from the specified column.

Parameters:

Name Type Description Default
fname str

Path to data file

 required
col int

Column index to plot (0-based)

 required
xlabel str

X-axis label, by default None

 None
ylabel str

Y-axis label, by default None

 None
**kwargs

Additional keyword arguments passed to ax.plot()

 {}

Returns:

Type Description
tuple

Tuple of (fig, ax, data) containing matplotlib figure, axes object, and loaded data
Source code in toolbox/plot/core.py 
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
def plot_lcurve(fname, col, xlabel=None, ylabel=None, **kwargs):
    """Plot learning curve from data file.

    This function reads a data file and plots a learning curve
    from the specified column.

    Parameters
    ----------
    fname : str
        Path to data file
    col : int
        Column index to plot (0-based)
    xlabel : str, optional
        X-axis label, by default None
    ylabel : str, optional
        Y-axis label, by default None
    **kwargs
        Additional keyword arguments passed to ax.plot()

    Returns
    -------
    tuple
        Tuple of (fig, ax, data) containing matplotlib figure,
        axes object, and loaded data
    """
    fig, ax = plt.subplots()
    data = np.loadtxt(fname)
    x = data[:, 0]
    ax.plot(x, data[:, col], **kwargs)
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(bottom=0.0)
    if xlabel is not None and ylabel is not None:
        ax_setlabel(ax, xlabel, ylabel)

    return fig, ax, data

plot_rmse(x, y, xlabel, ylabel, **kwargs)

Plot RMSE scatter plot with reference line.

This function creates a scatter plot of data points with a reference line and calculates RMSE.

Parameters:

Name Type Description Default
x array_like

X values

 required
y array_like

Y values

 required
xlabel str

X-axis label

 required
ylabel str

Y-axis label

 required
**kwargs

Additional keyword arguments passed to plotting functions

 {}

Returns:

Type Description
tuple

Tuple of (fig, ax, rmse) containing matplotlib figure, axes object, and calculated RMSE
Source code in toolbox/plot/core.py 
 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
117
118
119
120
121
122
123
124
125
126
def plot_rmse(x, y, xlabel, ylabel, **kwargs):
    """Plot RMSE scatter plot with reference line.

    This function creates a scatter plot of data points
    with a reference line and calculates RMSE.

    Parameters
    ----------
    x : array_like
        X values
    y : array_like
        Y values
    xlabel : str
        X-axis label
    ylabel : str
        Y-axis label
    **kwargs
        Additional keyword arguments passed to plotting functions

    Returns
    -------
    tuple
        Tuple of (fig, ax, rmse) containing matplotlib figure,
        axes object, and calculated RMSE
    """
    x = np.array(x)
    y = np.array(y)

    rmse = np.sqrt(metrics.mean_squared_error(x, y))

    fig, ax = plt.subplots(figsize=[4, 4])
    ax_rmse(ax, x, y)
    ax_setlabel(ax, xlabel, ylabel, **kwargs)

    return fig, ax, rmse

Deep Potential Plotting

toolbox.plot.dp

Deep Potential plotting module.

This module provides specialized plotting classes for visualizing Deep Potential (DP) training data and results.

DPTrainFigure

Bases: Figure

Figure class for DP training visualization.

This class extends Figure to provide specific functionality for visualizing Deep Potential training data.

Source code in toolbox/plot/dp.py 
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
class DPTrainFigure(Figure):
    """Figure class for DP training visualization.

    This class extends Figure to provide specific
    functionality for visualizing Deep Potential training data.
    """

    def __init__(self, ax) -> None:
        """Initialize DPTrainFigure.

        Parameters
        ----------
        ax : matplotlib.axes.Axes
            Matplotlib axes object
        """
        super().__init__(ax)

    def setup(self, x, y, **kwargs):
        """Set up DP training plot.

        Parameters
        ----------
        x : array_like
            Training data features
        y : array_like
            Training data targets
        **kwargs
            Additional keyword arguments
        """
        kwargs.update({"alpha": 0.2})
        self.ax.set_yscale("log")
        super().setup(x, y, **kwargs)
__init__(ax)

Initialize DPTrainFigure.

Parameters:

Name Type Description Default
ax Axes

Matplotlib axes object

 required
Source code in toolbox/plot/dp.py 
18
19
20
21
22
23
24
25
26
def __init__(self, ax) -> None:
    """Initialize DPTrainFigure.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Matplotlib axes object
    """
    super().__init__(ax)
setup(x, y, **kwargs)

Set up DP training plot.

Parameters:

Name Type Description Default
x array_like

Training data features

 required
y array_like

Training data targets

 required
**kwargs

Additional keyword arguments

 {}
Source code in toolbox/plot/dp.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def setup(self, x, y, **kwargs):
    """Set up DP training plot.

    Parameters
    ----------
    x : array_like
        Training data features
    y : array_like
        Training data targets
    **kwargs
        Additional keyword arguments
    """
    kwargs.update({"alpha": 0.2})
    self.ax.set_yscale("log")
    super().setup(x, y, **kwargs)

Figure Classes

toolbox.plot.figure

Figure classes module.

This module provides base and specialized figure classes for creating matplotlib plots with consistent styling and layout.

Figure

Base class for matplotlib figure creation.

This class provides a framework for creating matplotlib figures with consistent styling and layout.

Source code in toolbox/plot/figure.py 
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class Figure:
    """Base class for matplotlib figure creation.

    This class provides a framework for creating
    matplotlib figures with consistent styling and layout.
    """

    def __init__(self, **kwargs) -> None:
        """Initialize Figure.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments passed to plt.subplots
        """
        self.fig, self.ax = plt.subplots(**kwargs)

    def setup(self, x, y, xlim=None, ylim=None, **kwargs):
        """Set up basic plot with data.

        Parameters
        ----------
        x : array_like
            X data values
        y : array_like
            Y data values
        xlim : tuple, optional
            X-axis limits as (min, max), by default None
        ylim : tuple, optional
            Y-axis limits as (min, max), by default None
        **kwargs
            Additional keyword arguments passed to ax.plot
        """
        self.ax.plot(x, y, **kwargs)
        if xlim is None:
            xlim = (np.min(x), np.max(x))
        if ylim is None:
            ylim = (
                np.min(y) - (np.max(y) - np.min(y)) * 0.1,
                np.max(y) + (np.max(y) - np.min(y)) * 0.1,
            )
        self.xlim = xlim
        self.ylim = ylim
        self.ax.set_xlim(xlim)
        self.ax.set_ylim(ylim)

    def set_labels(self, kw, xlabel=None, ylabel=None, **kwargs):
        """Set axis labels from keyword dictionary.

        Parameters
        ----------
        kw : str
            Keyword to look up in label dictionary
        xlabel : str, optional
            X-axis label, by default None
        ylabel : str, optional
            Y-axis label, by default None
        **kwargs
            Additional keyword arguments passed to ax_setlabel
        """
        try:
            labels = label_dict.get(kw)
            xlabel = labels[0]
            ylabel = labels[1]
        except KeyError:
            assert (xlabel is not None) and (ylabel is not None)
            ax_setlabel(self.ax, xlabel, ylabel, **kwargs)
__init__(**kwargs)

Initialize Figure.

Parameters:

Name Type Description Default
**kwargs

Additional keyword arguments passed to plt.subplots

 {}
Source code in toolbox/plot/figure.py 
21
22
23
24
25
26
27
28
29
def __init__(self, **kwargs) -> None:
    """Initialize Figure.

    Parameters
    ----------
    **kwargs
        Additional keyword arguments passed to plt.subplots
    """
    self.fig, self.ax = plt.subplots(**kwargs)
set_labels(kw, xlabel=None, ylabel=None, **kwargs)

Set axis labels from keyword dictionary.

Parameters:

Name Type Description Default
kw str

Keyword to look up in label dictionary

 required
xlabel str

X-axis label, by default None

 None
ylabel str

Y-axis label, by default None

 None
**kwargs

Additional keyword arguments passed to ax_setlabel

 {}
Source code in toolbox/plot/figure.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def set_labels(self, kw, xlabel=None, ylabel=None, **kwargs):
    """Set axis labels from keyword dictionary.

    Parameters
    ----------
    kw : str
        Keyword to look up in label dictionary
    xlabel : str, optional
        X-axis label, by default None
    ylabel : str, optional
        Y-axis label, by default None
    **kwargs
        Additional keyword arguments passed to ax_setlabel
    """
    try:
        labels = label_dict.get(kw)
        xlabel = labels[0]
        ylabel = labels[1]
    except KeyError:
        assert (xlabel is not None) and (ylabel is not None)
        ax_setlabel(self.ax, xlabel, ylabel, **kwargs)
setup(x, y, xlim=None, ylim=None, **kwargs)

Set up basic plot with data.

Parameters:

Name Type Description Default
x array_like

X data values

 required
y array_like

Y data values

 required
xlim tuple

X-axis limits as (min, max), by default None

 None
ylim tuple

Y-axis limits as (min, max), by default None

 None
**kwargs

Additional keyword arguments passed to ax.plot

 {}
Source code in toolbox/plot/figure.py 
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
def setup(self, x, y, xlim=None, ylim=None, **kwargs):
    """Set up basic plot with data.

    Parameters
    ----------
    x : array_like
        X data values
    y : array_like
        Y data values
    xlim : tuple, optional
        X-axis limits as (min, max), by default None
    ylim : tuple, optional
        Y-axis limits as (min, max), by default None
    **kwargs
        Additional keyword arguments passed to ax.plot
    """
    self.ax.plot(x, y, **kwargs)
    if xlim is None:
        xlim = (np.min(x), np.max(x))
    if ylim is None:
        ylim = (
            np.min(y) - (np.max(y) - np.min(y)) * 0.1,
            np.max(y) + (np.max(y) - np.min(y)) * 0.1,
        )
    self.xlim = xlim
    self.ylim = ylim
    self.ax.set_xlim(xlim)
    self.ax.set_ylim(ylim)

FullCellFigure

Bases: Figure

Figure class for full cell visualization.

This class extends Figure to provide specific functionality for visualizing systems with full periodic boundary conditions.

Source code in toolbox/plot/figure.py 
 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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
class FullCellFigure(Figure):
    """Figure class for full cell visualization.

    This class extends Figure to provide specific
    functionality for visualizing systems with full periodic
    boundary conditions.
    """

    def __init__(self, **kwargs) -> None:
        """Initialize FullCellFigure.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments passed to parent class
        """
        super().__init__(**kwargs)

    def setup(self, x, y, xlim=None, ylim=None, z_surfs=None, **kwargs):
        """Set up full cell visualization plot.

        Parameters
        ----------
        x : array_like
            X data values
        y : array_like
            Y data values
        xlim : tuple, optional
            X-axis limits as (min, max), by default None
        ylim : tuple, optional
            Y-axis limits as (min, max), by default None
        z_surfs : list
            List of z-surface positions for visualization
        **kwargs
            Additional keyword arguments passed to parent setup method
        """
        super().setup(x, y, xlim, ylim, **kwargs)

        ax = self.ax
        ax.axvline(x=z_surfs[0])
        ax.axvline(x=z_surfs[1])
        ax.axvline(x=(z_surfs[0] + z_surfs[1]) / 2, ls="--")
        ax.fill_between(
            x,
            self.ylim[0],
            self.ylim[1],
            where=x <= z_surfs[0],
            facecolor="gray",
            alpha=0.5,
        )
        ax.fill_between(
            x,
            self.ylim[0],
            self.ylim[1],
            where=x >= z_surfs[1],
            facecolor="gray",
            alpha=0.5,
        )
__init__(**kwargs)

Initialize FullCellFigure.

Parameters:

Name Type Description Default
**kwargs

Additional keyword arguments passed to parent class

 {}
Source code in toolbox/plot/figure.py 
91
92
93
94
95
96
97
98
99
def __init__(self, **kwargs) -> None:
    """Initialize FullCellFigure.

    Parameters
    ----------
    **kwargs
        Additional keyword arguments passed to parent class
    """
    super().__init__(**kwargs)
setup(x, y, xlim=None, ylim=None, z_surfs=None, **kwargs)

Set up full cell visualization plot.

Parameters:

Name Type Description Default
x array_like

X data values

 required
y array_like

Y data values

 required
xlim tuple

X-axis limits as (min, max), by default None

 None
ylim tuple

Y-axis limits as (min, max), by default None

 None
z_surfs list

List of z-surface positions for visualization

 None
**kwargs

Additional keyword arguments passed to parent setup method

 {}
Source code in toolbox/plot/figure.py 
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
131
132
133
134
135
136
137
138
139
140
def setup(self, x, y, xlim=None, ylim=None, z_surfs=None, **kwargs):
    """Set up full cell visualization plot.

    Parameters
    ----------
    x : array_like
        X data values
    y : array_like
        Y data values
    xlim : tuple, optional
        X-axis limits as (min, max), by default None
    ylim : tuple, optional
        Y-axis limits as (min, max), by default None
    z_surfs : list
        List of z-surface positions for visualization
    **kwargs
        Additional keyword arguments passed to parent setup method
    """
    super().setup(x, y, xlim, ylim, **kwargs)

    ax = self.ax
    ax.axvline(x=z_surfs[0])
    ax.axvline(x=z_surfs[1])
    ax.axvline(x=(z_surfs[0] + z_surfs[1]) / 2, ls="--")
    ax.fill_between(
        x,
        self.ylim[0],
        self.ylim[1],
        where=x <= z_surfs[0],
        facecolor="gray",
        alpha=0.5,
    )
    ax.fill_between(
        x,
        self.ylim[0],
        self.ylim[1],
        where=x >= z_surfs[1],
        facecolor="gray",
        alpha=0.5,
    )

HalfCellFigure

Bases: Figure

Figure class for half cell visualization.

This class extends Figure to provide specific functionality for visualizing systems with half periodic boundary conditions (slab geometry).

Source code in toolbox/plot/figure.py 
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
class HalfCellFigure(Figure):
    """Figure class for half cell visualization.

    This class extends Figure to provide specific
    functionality for visualizing systems with half periodic
    boundary conditions (slab geometry).
    """

    def __init__(self, **kwargs) -> None:
        """Initialize HalfCellFigure.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments passed to parent class
        """
        super().__init__(**kwargs)

    def setup(self, x, y, xlim=None, ylim=None, z_surf=None, **kwargs):
        """Set up half cell visualization plot.

        Parameters
        ----------
        x : array_like
            X data values
        y : array_like
            Y data values
        xlim : tuple, optional
            X-axis limits as (min, max), by default None
        ylim : tuple, optional
            Y-axis limits as (min, max), by default None
        z_surf : float
            Z-surface position for visualization
        **kwargs
            Additional keyword arguments passed to parent setup method
        """
        super().setup(x, y, xlim, ylim, **kwargs)

        ax = self.ax
        ax.axvline(x=z_surf)
        ax.fill_between(
            x, self.ylim[0], self.ylim[1], where=x < z_surf, facecolor="gray", alpha=0.5
        )
__init__(**kwargs)

Initialize HalfCellFigure.

Parameters:

Name Type Description Default
**kwargs

Additional keyword arguments passed to parent class

 {}
Source code in toolbox/plot/figure.py 
151
152
153
154
155
156
157
158
159
def __init__(self, **kwargs) -> None:
    """Initialize HalfCellFigure.

    Parameters
    ----------
    **kwargs
        Additional keyword arguments passed to parent class
    """
    super().__init__(**kwargs)
setup(x, y, xlim=None, ylim=None, z_surf=None, **kwargs)

Set up half cell visualization plot.

Parameters:

Name Type Description Default
x array_like

X data values

 required
y array_like

Y data values

 required
xlim tuple

X-axis limits as (min, max), by default None

 None
ylim tuple

Y-axis limits as (min, max), by default None

 None
z_surf float

Z-surface position for visualization

 None
**kwargs

Additional keyword arguments passed to parent setup method

 {}
Source code in toolbox/plot/figure.py 
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
def setup(self, x, y, xlim=None, ylim=None, z_surf=None, **kwargs):
    """Set up half cell visualization plot.

    Parameters
    ----------
    x : array_like
        X data values
    y : array_like
        Y data values
    xlim : tuple, optional
        X-axis limits as (min, max), by default None
    ylim : tuple, optional
        Y-axis limits as (min, max), by default None
    z_surf : float
        Z-surface position for visualization
    **kwargs
        Additional keyword arguments passed to parent setup method
    """
    super().setup(x, y, xlim, ylim, **kwargs)

    ax = self.ax
    ax.axvline(x=z_surf)
    ax.fill_between(
        x, self.ylim[0], self.ylim[1], where=x < z_surf, facecolor="gray", alpha=0.5
    )

Statistical Plotting

toolbox.plot.stats

Statistical plotting module.

This module provides classes for creating statistical plots and performing statistical tests, including finite difference method testing.

FDMTest

Finite difference method test for derivative.

This class implements finite difference methods to test numerical derivatives against analytical solutions.

Source code in toolbox/plot/stats.py 
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
66
67
class FDMTest:
    """Finite difference method test for derivative.

    This class implements finite difference methods
    to test numerical derivatives against analytical solutions.
    """

    def __init__(self, **kwargs) -> None:
        """Initialize FDMTest.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments passed to plt.subplots
        """
        self.fig, self.axs = plt.subplots(
            nrows=2, figsize=[4, 6], sharex="all", **kwargs
        )

    def setup(self, x, y, dydx):
        """Set up finite difference test.

        Parameters
        ----------
        x : array_like
            X values
        y : array_like
            Y values
        dydx : array_like
            Analytical derivative values
        """
        self.stats_dict = {}
        ax = self.axs[0]
        ax.scatter(x, dydx, color="blue", label="original data")
        fit_output = stats.linregress(x, dydx)
        ax.plot(
            x,
            x * fit_output.slope + fit_output.intercept,
            "--",
            color="red",
            label="fitted line",
        )
        self.stats_dict["slope"] = fit_output.slope
        self.stats_dict["intercept"] = fit_output.intercept
        self.stats_dict["rvalue"] = fit_output.rvalue

        ax = self.axs[1]
        ref_data_diff = np.diff(y, axis=0).reshape(-1)
        test_data_diff = 0.5 * np.diff(x, axis=0) * (dydx[1:] + dydx[:-1])
        ax.scatter(x[1:], ref_data_diff, color="blue", label="original data")
        ax.plot(x[1:], test_data_diff, "--", color="red", label="fitted line")
        self.stats_dict["mean error"] = np.mean(ref_data_diff - test_data_diff)
__init__(**kwargs)

Initialize FDMTest.

Parameters:

Name Type Description Default
**kwargs

Additional keyword arguments passed to plt.subplots

 {}
Source code in toolbox/plot/stats.py 
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, **kwargs) -> None:
    """Initialize FDMTest.

    Parameters
    ----------
    **kwargs
        Additional keyword arguments passed to plt.subplots
    """
    self.fig, self.axs = plt.subplots(
        nrows=2, figsize=[4, 6], sharex="all", **kwargs
    )
setup(x, y, dydx)

Set up finite difference test.

Parameters:

Name Type Description Default
x array_like

X values

 required
y array_like

Y values

 required
dydx array_like

Analytical derivative values

 required
Source code in toolbox/plot/stats.py 
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
def setup(self, x, y, dydx):
    """Set up finite difference test.

    Parameters
    ----------
    x : array_like
        X values
    y : array_like
        Y values
    dydx : array_like
        Analytical derivative values
    """
    self.stats_dict = {}
    ax = self.axs[0]
    ax.scatter(x, dydx, color="blue", label="original data")
    fit_output = stats.linregress(x, dydx)
    ax.plot(
        x,
        x * fit_output.slope + fit_output.intercept,
        "--",
        color="red",
        label="fitted line",
    )
    self.stats_dict["slope"] = fit_output.slope
    self.stats_dict["intercept"] = fit_output.intercept
    self.stats_dict["rvalue"] = fit_output.rvalue

    ax = self.axs[1]
    ref_data_diff = np.diff(y, axis=0).reshape(-1)
    test_data_diff = 0.5 * np.diff(x, axis=0) * (dydx[1:] + dydx[:-1])
    ax.scatter(x[1:], ref_data_diff, color="blue", label="original data")
    ax.plot(x[1:], test_data_diff, "--", color="red", label="fitted line")
    self.stats_dict["mean error"] = np.mean(ref_data_diff - test_data_diff)

Plot Styles

toolbox.plot.style

Matplotlib style module.

This module provides functionality for using custom matplotlib styles and color maps for scientific plotting.

References

use_style(style_name)

Use custom matplotlib style.

Parameters:

Name Type Description Default
style_name str

Name of the style to use

 required
Notes

If the specified style is not found, falls back to fivethirtyeight style and prints a warning.

Source code in toolbox/plot/style.py 
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
def use_style(style_name):
    """Use custom matplotlib style.

    Parameters
    ----------
    style_name : str
        Name of the style to use

    Notes
    -----
    If the specified style is not found, falls back to fivethirtyeight
    style and prints a warning.
    """
    plt.style.use("fivethirtyeight")
    fname = str(MODULE_DIR / (f"mplstyle/{style_name}.mplstyle"))
    # print(fname)
    try:
        plt.style.use(fname)
    except OSError:
        print(
            f"Warning: no style {style_name} is found. Use matplotlib default style."
        )
    """
    Set colors:
    https://stackoverflow.com/questions/68664116/is-there-a-way-to-change-the-color-names-color-in-mplstyle-file
    """
    color_map = colors.get_named_colors_mapping()
    color_map["red"] = "#D7422A"
    color_map["blue"] = "#003C88"

Utils Module

Bibliography Tools

toolbox.utils.bibtool

Bibliography tool module.

This module provides utilities for managing bibliography files, including extracting citations from LaTeX files and exporting selected bibliography entries.

export(bib_in_file, tex_files=None, bib_out_file='export-ref.bib', online=True)

Export a subset of the bib file that is used in the tex file.

Parameters:

Name Type Description Default
bib_in_file str

Path to the original bib file

 required
tex_files Union[List[str], str]

Path to the tex file or list of tex files

 None
bib_out_file str

Path to the output bib file

 'export-ref.bib'
online bool

If True, it will try to get the bib entry from the DOI

 True
Source code in toolbox/utils/bibtool.py 
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
def export(
    bib_in_file: str,
    tex_files: Optional[Union[List[str], str]] = None,
    bib_out_file: str = "export-ref.bib",
    online: bool = True,
):
    """
    Export a subset of the bib file that is used in the tex file.

    Parameters
    ----------
    bib_in_file : str
        Path to the original bib file
    tex_files : Union[List[str], str]
        Path to the tex file or list of tex files
    bib_out_file : str
        Path to the output bib file
    online : bool
        If True, it will try to get the bib entry from the DOI
    """
    if tex_files is None:
        tex_files = glob.glob("./*.tex")
    if isinstance(tex_files, str):
        tex_files = [tex_files]

    citation_keys = extract_citation_keys(tex_files)
    bib_data = parse_file(bib_in_file)

    new_bib_data = BibliographyData()
    for kw in tqdm(citation_keys):
        if kw in bib_data.entries:
            if online:
                try:
                    out = get_bib_from_doi(
                        bib_data.entries[kw].fields["doi"], abbrev_journal=True
                    )
                    obj = parse_string(out[1], bib_format="bibtex")
                    for tmp_kw in obj.entries:
                        print(tmp_kw)
                    new_bib_data.entries[kw] = obj.entries[tmp_kw]
                except KeyError:
                    new_bib_data.entries[kw] = bib_data.entries[kw]
                new_bib_data.entries[kw].fields["title"] = bib_data.entries[kw].fields[
                    "title"
                ]
                with contextlib.suppress(KeyError):
                    new_bib_data.entries[kw].fields["journal"] = bib_data.entries[
                        kw
                    ].fields["journal"]
            else:
                new_bib_data.entries[kw] = bib_data.entries[kw]

    with open(bib_out_file, "w") as new_bib_file:
        new_bib_file.write(new_bib_data.to_string("bibtex"))

extract_citation_keys(fnames)

Grep all citation keys used in the tex file.

Parameters:

Name Type Description Default
fnames str

Path to the tex file

 required
Source code in toolbox/utils/bibtool.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def extract_citation_keys(fnames: list[str]):
    """
    Grep all citation keys used in the tex file.

    Parameters
    ----------
    fnames : str
        Path to the tex file
    """
    citation_keys = set()
    for fname in fnames:
        with open(fname) as file:
            tex_content = file.read()
        _citation_keys = set(re.findall(r"\\cite{([^}]+)}", tex_content))
        _citation_keys = {
            key.strip() for keys in _citation_keys for key in keys.split(",")
        }
        citation_keys.update(_citation_keys)
    return citation_keys

Data Utilities

toolbox.utils.data

Data system management module.

This module provides base classes for managing data system configurations with read and write functionality.

DataSystem

Base class for data system management.

This class provides a framework for reading and writing data system configurations.

Source code in toolbox/utils/data.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class DataSystem:
    """Base class for data system management.

    This class provides a framework for reading and writing
    data system configurations.
    """

    def __init__(self) -> None:
        """Initialize DataSystem."""
        pass

    def read(self):
        """Read data system configuration."""
        pass

    def write(self):
        """Write data system configuration."""
        pass
__init__()

Initialize DataSystem.

Source code in toolbox/utils/data.py 
15
16
17
def __init__(self) -> None:
    """Initialize DataSystem."""
    pass
read()

Read data system configuration.

Source code in toolbox/utils/data.py 
19
20
21
def read(self):
    """Read data system configuration."""
    pass
write()

Write data system configuration.

Source code in toolbox/utils/data.py 
23
24
25
def write(self):
    """Write data system configuration."""
    pass

Mathematical Functions

toolbox.utils.math

Mathematical utilities module.

This module provides various mathematical functions for data analysis, including statistical functions, interpolation, integration, filtering, and error metrics.

block_ave(_x, _y, l_block)

Calculate block averages of data.

Parameters:

Name Type Description Default
_x array_like

Input x values

 required
_y array_like

Input y values

 required
l_block int

Block size for averaging

 required

Returns:

Type Description
tuple

Tuple of (x_ave, y_ave) containing block-averaged values
Source code in toolbox/utils/math.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def block_ave(_x, _y, l_block):
    """Calculate block averages of data.

    Parameters
    ----------
    _x : array_like
        Input x values
    _y : array_like
        Input y values
    l_block : int
        Block size for averaging

    Returns
    -------
    tuple
        Tuple of (x_ave, y_ave) containing block-averaged values
    """
    assert len(_x) == len(_y)
    n_block = math.floor(len(_x) / l_block)
    x = _x[: (n_block * l_block)]
    y = _y[: (n_block * l_block)]
    x = np.reshape(x, (-1, l_block)).mean(axis=-1)
    y = np.reshape(y, (-1, l_block)).mean(axis=-1)
    return x, y

cumave(data)

Calculate cumulative average of data.

Parameters:

Name Type Description Default
data array_like

Input data

 required

Returns:

Type Description
array_like

Cumulative average values
Source code in toolbox/utils/math.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def cumave(data):
    """Calculate cumulative average of data.

    Parameters
    ----------
    data : array_like
        Input data

    Returns
    -------
    array_like
        Cumulative average values
    """
    cum_sum = data.cumsum()
    cum_ave = cum_sum / (np.arange(len(data)) + 1)
    return cum_ave

error_test(y_true, y_pred)

Calculate regression error metrics.

This function calculates various regression error metrics between true and predicted values.

Parameters:

Name Type Description Default
y_true array_like

True values

 required
y_pred array_like

Predicted values

 required

Returns:

Type Description
dict

Dictionary containing error metrics: - max_err: Maximum error - mae: Mean absolute error - rmse: Root mean squared error - r2: R-squared score - mape: Mean absolute percentage error
References

https://scikit-learn.org/stable/modules/model_evaluation.html#regression-metrics

Source code in toolbox/utils/math.py 
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
def error_test(y_true, y_pred):
    """Calculate regression error metrics.

    This function calculates various regression error metrics
    between true and predicted values.

    Parameters
    ----------
    y_true : array_like
        True values
    y_pred : array_like
        Predicted values

    Returns
    -------
    dict
        Dictionary containing error metrics:
        - max_err: Maximum error
        - mae: Mean absolute error
        - rmse: Root mean squared error
        - r2: R-squared score
        - mape: Mean absolute percentage error

    References
    ----------
    https://scikit-learn.org/stable/modules/model_evaluation.html#regression-metrics
    """
    results_dict = {}
    y_true = np.reshape(y_true, (-1,))
    y_pred = np.reshape(y_pred, (-1,))

    results_dict["max_err"] = metrics.max_error(y_true, y_pred)
    results_dict["mae"] = metrics.mean_absolute_error(y_true, y_pred)
    results_dict["rmse"] = np.sqrt(metrics.mean_squared_error(y_true, y_pred))
    results_dict["r2"] = metrics.r2_score(y_true, y_pred)
    results_dict["mape"] = metrics.mean_absolute_percentage_error(y_true, y_pred)
    return results_dict

gaussian_filter(data, bin_edge, sigma, weight=None)

Apply Gaussian filter to data.

Parameters:

Name Type Description Default
data array_like

Input data to filter

 required
bin_edge array_like

Bin edges for filtering

 required
sigma float

Gaussian standard deviation

 required
weight array_like

Weights for each data point

 None

Returns:

Type Description
tuple

Tuple of (bins, filtered_data)
Source code in toolbox/utils/math.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
def gaussian_filter(data, bin_edge, sigma: float, weight=None):
    """Apply Gaussian filter to data.

    Parameters
    ----------
    data : array_like
        Input data to filter
    bin_edge : array_like
        Bin edges for filtering
    sigma : float
        Gaussian standard deviation
    weight : array_like, optional
        Weights for each data point

    Returns
    -------
    tuple
        Tuple of (bins, filtered_data)
    """
    data = np.reshape(data, (-1, 1))
    bins = get_bins_from_bin_edge(bin_edge)
    bins = np.reshape(bins, (1, -1))
    weight = np.ones_like(data) if weight is None else np.reshape(weight, (-1, 1))
    output = (
        np.exp(-(((bins - data) / sigma) ** 2)) / (np.sqrt(2 * np.pi) * sigma) * weight
    )
    return bins, output.sum(axis=0)

gaussian_func(x, mu=0, sigma=1)

Gaussian probability density function.

Parameters:

Name Type Description Default
x array_like

Input values

 required
mu float

Mean of the distribution, by default 0

 0
sigma float

Standard deviation, by default 1

 1

Returns:

Type Description
array_like

Gaussian probability density values
Source code in toolbox/utils/math.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def gaussian_func(x, mu=0, sigma=1):
    """Gaussian probability density function.

    Parameters
    ----------
    x : array_like
        Input values
    mu : float, optional
        Mean of the distribution, by default 0
    sigma : float, optional
        Standard deviation, by default 1

    Returns
    -------
    array_like
        Gaussian probability density values
    """
    return stats.norm.pdf(x, loc=mu, scale=sigma)

gaussian_int(x, mu, sigma)

Gaussian cumulative distribution function.

Parameters:

Name Type Description Default
x array_like

Input values

 required
mu float

Mean of the distribution

 required
sigma float

Standard deviation

 required

Returns:

Type Description
array_like

Gaussian cumulative distribution values
Source code in toolbox/utils/math.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def gaussian_int(x, mu, sigma):
    """Gaussian cumulative distribution function.

    Parameters
    ----------
    x : array_like
        Input values
    mu : float
        Mean of the distribution
    sigma : float
        Standard deviation

    Returns
    -------
    array_like
        Gaussian cumulative distribution values
    """
    coeff = 1 / 2
    return coeff * erf((x - mu) / (np.sqrt(2) * sigma))

get_bin_ave(data, bin_size=10)

Calculate binned averages of data.

Parameters:

Name Type Description Default
data array_like

Input data to bin

 required
bin_size int

Size of each bin, by default 10

 10

Returns:

Type Description
array_like

Binned averages
Source code in toolbox/utils/math.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def get_bin_ave(data, bin_size=10):
    """Calculate binned averages of data.

    Parameters
    ----------
    data : array_like
        Input data to bin
    bin_size : int, optional
        Size of each bin, by default 10

    Returns
    -------
    array_like
        Binned averages
    """
    total_size = len(data)
    out_size = bin_size * total_size // bin_size
    data = np.reshape(data[:out_size], (-1, bin_size))
    data = np.mean(data, axis=-1)
    return data

get_dev(x, y)

Calculate derivative dy/dx.

Parameters:

Name Type Description Default
x array_like

Input x values

 required
y array_like

Input y values

 required

Returns:

Type Description
tuple

Tuple of (x_mid, dy/dx) where x_mid are midpoints
Source code in toolbox/utils/math.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def get_dev(x, y):
    """Calculate derivative dy/dx.

    Parameters
    ----------
    x : array_like
        Input x values
    y : array_like
        Input y values

    Returns
    -------
    tuple
        Tuple of (x_mid, dy/dx) where x_mid are midpoints
    """
    x = np.array(x)
    y = np.array(y)
    delta_x = np.diff(x)
    delta_y = np.diff(y)
    out_x = (x[1:] + x[:-1]) / 2
    return out_x, delta_y / delta_x

get_int(x, y)

Calculate numerical integral using trapezoidal rule.

Parameters:

Name Type Description Default
x array_like

Input x values

 required
y array_like

Input y values

 required

Returns:

Type Description
tuple

Tuple of (x_mid, integral) where x_mid are midpoints
Source code in toolbox/utils/math.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def get_int(x, y):
    """Calculate numerical integral using trapezoidal rule.

    Parameters
    ----------
    x : array_like
        Input x values
    y : array_like
        Input y values

    Returns
    -------
    tuple
        Tuple of (x_mid, integral) where x_mid are midpoints
    """
    x = np.array(x)
    y = np.array(y)
    ave_y = (y[1:] + y[:-1]) / 2
    delta_x = np.diff(x)
    x = (x[1:] + x[:-1]) / 2
    return x, np.cumsum(ave_y * delta_x)

get_int_array(x, ys)

Calculate numerical integrals for multiple y arrays.

Parameters:

Name Type Description Default
x array_like

Input x values

 required
ys array_like

Multiple y arrays with shape (n_arrays, n_points)

 required

Returns:

Type Description
tuple

Tuple of (x_mid, integrals) where x_mid are midpoints
Source code in toolbox/utils/math.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def get_int_array(x, ys):
    """Calculate numerical integrals for multiple y arrays.

    Parameters
    ----------
    x : array_like
        Input x values
    ys : array_like
        Multiple y arrays with shape (n_arrays, n_points)

    Returns
    -------
    tuple
        Tuple of (x_mid, integrals) where x_mid are midpoints
    """
    x = np.array(x)
    ys = np.array(ys)
    ave_y = (ys[:, 1:] + ys[:, :-1]) / 2
    delta_x = np.diff(x)
    x = (x[1:] + x[:-1]) / 2
    return x, np.cumsum(ave_y * delta_x.reshape(1, -1), axis=-1)

handle_zero_division(x, y, threshold=None)

Handle division by zero with optional threshold.

This function safely divides x by y, handling zero division with optional threshold masking.

Parameters:

Name Type Description Default
x array_like

Numerator values

 required
y array_like

Denominator values

 required
threshold float

Threshold below which y values are set to zero

 None

Returns:

Type Description
array_like

Result of x/y with safe division
Source code in toolbox/utils/math.py 
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
def handle_zero_division(x, y, threshold=None):
    """Handle division by zero with optional threshold.

    This function safely divides x by y, handling zero
    division with optional threshold masking.

    Parameters
    ----------
    x : array_like
        Numerator values
    y : array_like
        Denominator values
    threshold : float, optional
        Threshold below which y values are set to zero

    Returns
    -------
    array_like
        Result of x/y with safe division
    """
    if threshold is not None:
        mask = np.abs(y) <= threshold
        y[np.nonzero(mask)] = 0.0
    with np.errstate(divide="ignore", invalid="ignore"):
        result = np.true_divide(x, y)
        # replace NaN and Inf values with 0
        result[~np.isfinite(result)] = 0
    return result

interp(x, dataset)

Interpolate y values from dataset.

Parameters:

Name Type Description Default
x array_like

X values to interpolate at

 required
dataset array_like

Dataset with dataset[0] = x_values, dataset[1] = y_values

 required

Returns:

Type Description
array_like

Interpolated y values
Source code in toolbox/utils/math.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def interp(x, dataset):
    """Interpolate y values from dataset.

    Parameters
    ----------
    x : array_like
        X values to interpolate at
    dataset : array_like
        Dataset with dataset[0] = x_values, dataset[1] = y_values

    Returns
    -------
    array_like
        Interpolated y values
    """
    y = np.interp(x, xp=dataset[0], fp=dataset[1])
    return y

vec_project(vec, unit_vec)

Project vector onto unit vector.

Parameters:

Name Type Description Default
vec array_like

Vector to project

 required
unit_vec array_like

Unit vector to project onto

 required

Returns:

Type Description
array_like

Projected vector
Source code in toolbox/utils/math.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def vec_project(vec, unit_vec):
    """Project vector onto unit vector.

    Parameters
    ----------
    vec : array_like
        Vector to project
    unit_vec : array_like
        Unit vector to project onto

    Returns
    -------
    array_like
        Projected vector
    """
    return np.dot(vec, unit_vec) * unit_vec

Optimizer

toolbox.utils.optimizer

Optimizer module.

This module provides classes and functions for optimization using various linear regression methods.

Optimizer

Optimizer for linear regression models.

This class provides various optimization methods for fitting linear models to data.

Source code in toolbox/utils/optimizer.py 
 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
 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class Optimizer:
    """Optimizer for linear regression models.

    This class provides various optimization methods for
    fitting linear models to data.
    """

    def __init__(self, method="ols_cut") -> None:
        """Initialize Optimizer.

        Parameters
        ----------
        method : str, optional
            Optimization method to use, by default "ols_cut"
        """
        self.method = method

    def run(self, x, y, **kwargs):
        """Run optimization with specified method.

        Parameters
        ----------
        x : array_like
            Input features
        y : array_like
            Target values
        **kwargs
            Additional keyword arguments

        Returns
        -------
        array_like
            Optimization result
        """
        assert len(x) == len(y)
        output = getattr(self, f"_run_{self.method}")(x, y, **kwargs)
        return output

    def _run_ols(self, x, y):
        """Run ordinary least squares optimization.

        Parameters
        ----------
        x : array_like
            Input features
        y : array_like
            Target values

        Returns
        -------
        array_like
            Optimization result [-intercept/slope]
        """
        result = stats.linregress(x=x, y=y)
        output = -result.intercept / result.slope
        return output

    def _run_ols_cut(self, x, y, nstep=4):
        """Run OLS with cutoff on recent data.

        Parameters
        ----------
        x : array_like
            Input features
        y : array_like
            Target values
        nstep : int, optional
            Number of recent steps to use, by default 4

        Returns
        -------
        array_like
            Optimization result [-intercept/slope]
        """
        l_cut = min(len(x), nstep)
        result = stats.linregress(x=x[-l_cut:], y=y[-l_cut:])
        output = -result.intercept / result.slope
        return output

    def _run_wls(self, X, y):
        """Run weighted least squares optimization.

        Parameters
        ----------
        X : array_like
            Input features with constant term
        y : array_like
            Target values

        Returns
        -------
        array_like
            Optimization result [-intercept/slope]
        """
        # fit linear regression model
        X = sm.add_constant(self.x)
        wt = np.exp(-(np.array(self.y) ** 2) / 0.1)
        fit_wls = sm.WLS(self.y, X, weights=wt).fit()
        return -fit_wls.params[0] / fit_wls.params[1]

    def _run_wls_cut(self, X, y, nstep=4):
        """Run WLS with cutoff on recent data.

        Parameters
        ----------
        X : array_like
            Input features with constant term
        y : array_like
            Target values
        nstep : int, optional
            Number of recent steps to use, by default 4

        Returns
        -------
        array_like
            Optimization result [-intercept/slope]
        """
        l_cut = min(len(X), nstep)
        X = X[-l_cut:]
        y = y[-l_cut:]
        X = sm.add_constant(X)
        wt = np.exp(-(np.array(y) ** 2) / 0.1)
        fit_wls = sm.WLS(y, X, weights=wt).fit()
        return -fit_wls.params[0] / fit_wls.params[1]
__init__(method='ols_cut')

Initialize Optimizer.

Parameters:

Name Type Description Default
method str

Optimization method to use, by default "ols_cut"

 'ols_cut'
Source code in toolbox/utils/optimizer.py 
20
21
22
23
24
25
26
27
28
def __init__(self, method="ols_cut") -> None:
    """Initialize Optimizer.

    Parameters
    ----------
    method : str, optional
        Optimization method to use, by default "ols_cut"
    """
    self.method = method
run(x, y, **kwargs)

Run optimization with specified method.

Parameters:

Name Type Description Default
x array_like

Input features

 required
y array_like

Target values

 required
**kwargs

Additional keyword arguments

 {}

Returns:

Type Description
array_like

Optimization result
Source code in toolbox/utils/optimizer.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def run(self, x, y, **kwargs):
    """Run optimization with specified method.

    Parameters
    ----------
    x : array_like
        Input features
    y : array_like
        Target values
    **kwargs
        Additional keyword arguments

    Returns
    -------
    array_like
        Optimization result
    """
    assert len(x) == len(y)
    output = getattr(self, f"_run_{self.method}")(x, y, **kwargs)
    return output

Toy Models

toolbox.utils.toy_model

Toy model module.

This module provides simple physics models for educational purposes, including a parallel plate capacitor model.

ParallelPlateCapacitor

Class for parallel plate capacitor calculations.

This class models a parallel plate capacitor system with configurable permittivity, plate distance, and capacitance.

Parameters:

Name Type Description Default
epsilon_r float

Relative permittivity

 None
d float

Distance between plates in Angstroms

 None
capacitance float

Capacitance in μF/cm²

 None

Examples:

>>> obj = ParallelPlateCapacitor(epsilon_r=10.0, d=4.0)
>>> print(obj)
>>> obj = ParallelPlateCapacitor(capacitance=20.0, d=4.0)
>>> print(obj)
Source code in toolbox/utils/toy_model.py 
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
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
class ParallelPlateCapacitor:
    """Class for parallel plate capacitor calculations.

    This class models a parallel plate capacitor system
    with configurable permittivity, plate distance, and capacitance.

    Parameters
    ----------
    epsilon_r : float
        Relative permittivity
    d : float
        Distance between plates in Angstroms
    capacitance : float
        Capacitance in μF/cm²

    Examples
    --------
    >>> obj = ParallelPlateCapacitor(epsilon_r=10.0, d=4.0)
    >>> print(obj)
    >>> obj = ParallelPlateCapacitor(capacitance=20.0, d=4.0)
    >>> print(obj)
    """

    def __init__(
        self,
        epsilon_r: Optional[float] = None,
        d: Optional[float] = None,
        capacitance: Optional[float] = None,
    ):
        """Initialize ParallelPlateCapacitor.

        Parameters
        ----------
        epsilon_r : Optional[float], optional
            Relative permittivity, by default None
        d : Optional[float], optional
            Distance between plates in Angstroms, by default None
        capacitance : Optional[float], optional
            Capacitance in μF/cm², by default None
        """
        # conversion factor from F/m² to μF/cm²
        self._coeff = constants.centi**2 / constants.micro

        if epsilon_r is not None:
            self._epsilon_r = epsilon_r

        if d is not None:
            # d [m]
            self._d = d * constants.angstrom

        if capacitance is not None:
            self._capacitance = capacitance

    @property
    def capacitance(self):
        """Get capacitance in μF/cm²."""
        try:
            return self._capacitance
        except AttributeError:
            return self._epsilon_r * constants.epsilon_0 / self._d * self._coeff

    @property
    def epsilon_r(self):
        """Get relative permittivity."""
        try:
            return self._epsilon_r
        except AttributeError:
            return self._capacitance * self._d / (constants.epsilon_0 * self._coeff)

    @property
    def d_angstrom(self):
        """Get plate distance in Angstroms."""
        return self._d / constants.angstrom

    def __repr__(self):
        """Return string representation of the capacitor.

        Returns
        -------
        str
            String representation showing key parameters
        """
        return f"ParallelPlateCapacitor(epsilon_r={self.epsilon_r}, d={self.d_angstrom} angstrom, capacitance={self.capacitance} muF/cm^2)"
capacitance property

Get capacitance in μF/cm².

d_angstrom property

Get plate distance in Angstroms.

epsilon_r property

Get relative permittivity.

__init__(epsilon_r=None, d=None, capacitance=None)

Initialize ParallelPlateCapacitor.

Parameters:

Name Type Description Default
epsilon_r Optional[float]

Relative permittivity, by default None

 None
d Optional[float]

Distance between plates in Angstroms, by default None

 None
capacitance Optional[float]

Capacitance in μF/cm², by default None

 None
Source code in toolbox/utils/toy_model.py 
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
def __init__(
    self,
    epsilon_r: Optional[float] = None,
    d: Optional[float] = None,
    capacitance: Optional[float] = None,
):
    """Initialize ParallelPlateCapacitor.

    Parameters
    ----------
    epsilon_r : Optional[float], optional
        Relative permittivity, by default None
    d : Optional[float], optional
        Distance between plates in Angstroms, by default None
    capacitance : Optional[float], optional
        Capacitance in μF/cm², by default None
    """
    # conversion factor from F/m² to μF/cm²
    self._coeff = constants.centi**2 / constants.micro

    if epsilon_r is not None:
        self._epsilon_r = epsilon_r

    if d is not None:
        # d [m]
        self._d = d * constants.angstrom

    if capacitance is not None:
        self._capacitance = capacitance
__repr__()

Return string representation of the capacitor.

Returns:

Type Description
str

String representation showing key parameters
Source code in toolbox/utils/toy_model.py 
87
88
89
90
91
92
93
94
95
def __repr__(self):
    """Return string representation of the capacitor.

    Returns
    -------
    str
        String representation showing key parameters
    """
    return f"ParallelPlateCapacitor(epsilon_r={self.epsilon_r}, d={self.d_angstrom} angstrom, capacitance={self.capacitance} muF/cm^2)"

Unit Conversion

toolbox.utils.unit

Physical constants and unit conversions.

References

https://docs.scipy.org/doc/scipy/reference/constants.html

General Utilities

toolbox.utils.utils

Utility functions for computational chemistry and materials science.

This module provides a collection of utility functions for various tasks including: - Dictionary manipulation and file I/O operations - CP2K input file generation - Density and concentration calculations - Water molecule analysis and manipulation - Lennard-Jones parameter calculations - Coordinate number calculations

The functions are designed to work with common scientific computing libraries such as NumPy, ASE, and MDAnalysis.

calc_coord_number(atoms, c_ids, neigh_ids, cutoff=None, voronoi=False)

Calculate coordination numbers for specified atoms.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object

 required
c_ids ndarray

Indices of central atoms

 required
neigh_ids ndarray

Indices of neighbor atoms

 required
cutoff Optional[float]

Distance cutoff for coordination calculation

 None
voronoi bool

If True, use Voronoi method instead of cutoff

 False

Returns:

Type Description
ndarray

Array of coordination numbers
Source code in toolbox/utils/utils.py 
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
def calc_coord_number(
    atoms: Atoms,
    c_ids: np.ndarray,
    neigh_ids: np.ndarray,
    cutoff: Optional[float] = None,
    voronoi: bool = False,
) -> np.ndarray:
    """
    Calculate coordination numbers for specified atoms.

    Parameters
    ----------
    atoms : Atoms
        ASE Atoms object
    c_ids : np.ndarray
        Indices of central atoms
    neigh_ids : np.ndarray
        Indices of neighbor atoms
    cutoff : Optional[float], optional
        Distance cutoff for coordination calculation
    voronoi : bool, optional
        If True, use Voronoi method instead of cutoff

    Returns
    -------
    np.ndarray
        Array of coordination numbers
    """
    p = atoms.get_positions()
    p_c = p[c_ids]
    p_n = p[neigh_ids]
    results = np.empty((len(c_ids), len(neigh_ids)), dtype=np.float64)
    distance_array(p_c, p_n, box=atoms.cell.cellpar(), result=results)
    if cutoff is None and voronoi:
        # use voronoi method
        out = np.unique(np.argmin(results, axis=0), return_counts=True)
        cns = np.zeros(len(p_c), dtype=np.int32)
        cns[out[0]] = out[1]
    else:
        # use cutoff method
        cns = np.count_nonzero(results <= cutoff, axis=1)
    return cns

calc_density(n, v, mol_mass)

Calculate density (g/cm³) from the number of particles.

Parameters:

Name Type Description Default
n int or ndarray

Number of particles

 required
v float or ndarray

Volume in ų

 required
mol_mass float

Molar mass in g/mol

 required

Returns:

Type Description
float or ndarray

Density in g/cm³
Source code in toolbox/utils/utils.py 
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
def calc_density(
    n: Union[int, np.ndarray], v: Union[float, np.ndarray], mol_mass: float
) -> Union[float, np.ndarray]:
    """
    Calculate density (g/cm³) from the number of particles.

    Parameters
    ----------
    n : int or np.ndarray
        Number of particles
    v : float or np.ndarray
        Volume in ų
    mol_mass : float
        Molar mass in g/mol

    Returns
    -------
    float or np.ndarray
        Density in g/cm³
    """
    rho = (n / constants.Avogadro * mol_mass) / (
        v * (constants.angstrom / constants.centi) ** 3
    )
    return rho

calc_lj_params(ks)

Calculate and print Lennard-Jones parameters for element pairs.

This function calculates mixed Lennard-Jones parameters using the Lorentz-Berthelot combining rules and prints them to stdout.

Parameters:

Name Type Description Default
ks List[str]

List of element symbols

 required
Source code in toolbox/utils/utils.py 
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
def calc_lj_params(ks: list[str]) -> None:
    """
    Calculate and print Lennard-Jones parameters for element pairs.

    This function calculates mixed Lennard-Jones parameters using the
    Lorentz-Berthelot combining rules and prints them to stdout.

    Parameters
    ----------
    ks : List[str]
        List of element symbols
    """
    n_elements = len(ks)
    _sigma = []
    _epsilon = []
    for i in range(n_elements):
        try:
            sigma_i = lj_params[ks[i]]["sigma"]
            epsilon_i = lj_params[ks[i]]["epsilon"]
            _sigma.append(sigma_i)
        except KeyError as exc:
            raise KeyError(f"sigma for {ks[i]} not found") from exc
        for j in range(i, n_elements):
            try:
                sigma_j = lj_params[ks[j]]["sigma"]
                epsilon_j = lj_params[ks[j]]["epsilon"]
            except KeyError as exc:
                raise KeyError(f"sigma for {ks[j]} not found") from exc
            print(ks[i], ks[j], (sigma_i + sigma_j) / 2, np.sqrt(epsilon_i * epsilon_j))

calc_molar_concentration(number_density, grid_volume)

Calculate molar concentration from number density.

Parameters:

Name Type Description Default
number_density int, float, or np.ndarray

Number of particles per grid

 required
grid_volume float or ndarray

Volume of each grid in ų (angstrom cubed)

 required

Returns:

Type Description
float or ndarray

Molar concentration in mol/L (moles of ions per liter of solution)
Source code in toolbox/utils/utils.py 
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
def calc_molar_concentration(
    number_density: Union[int, float, np.ndarray], grid_volume: Union[float, np.ndarray]
) -> Union[float, np.ndarray]:
    """
    Calculate molar concentration from number density.

    Parameters
    ----------
    number_density : int, float, or np.ndarray
        Number of particles per grid
    grid_volume : float or np.ndarray
        Volume of each grid in ų (angstrom cubed)

    Returns
    -------
    float or np.ndarray
        Molar concentration in mol/L (moles of ions per liter of solution)
    """
    # Convert grid volume from ų to cm³ (1 ų = 1e-24 cm³)
    volume_cm3 = grid_volume * 1e-24

    # Convert volume from cm³ to liters (1 L = 1000 cm³)
    volume_liters = volume_cm3 / 1000.0

    # Calculate moles of ions in each grid
    moles = number_density / constants.Avogadro

    # Calculate molar concentration (mol/L)
    molar_concentration = moles / volume_liters

    return molar_concentration

calc_number(rho, v, mol_mass)

Calculate number of particles from density and volume.

Parameters:

Name Type Description Default
rho float or ndarray

Density in g/cm³

 required
v float or ndarray

Volume in ų

 required
mol_mass float

Molar mass in g/mol

 required

Returns:

Type Description
int

Number of particles (rounded to integer)
Source code in toolbox/utils/utils.py 
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
def calc_number(rho: Union[float, np.ndarray], v: Union[float, np.ndarray], mol_mass: float) -> int:
    """
    Calculate number of particles from density and volume.

    Parameters
    ----------
    rho : float or np.ndarray
        Density in g/cm³
    v : float or np.ndarray
        Volume in ų
    mol_mass : float
        Molar mass in g/mol

    Returns
    -------
    int
        Number of particles (rounded to integer)
    """
    n = (
        rho
        * (v * (constants.angstrom / constants.centi) ** 3)
        * constants.Avogadro
        / mol_mass
    )
    return int(n)

calc_water_coord_number(atoms)

Calculate coordination numbers for water molecules (O-H coordination).

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object containing water molecules

 required

Returns:

Type Description
ndarray

Array of coordination numbers for oxygen atoms
Source code in toolbox/utils/utils.py 
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
def calc_water_coord_number(atoms: Atoms) -> np.ndarray:
    """
    Calculate coordination numbers for water molecules (O-H coordination).

    Parameters
    ----------
    atoms : Atoms
        ASE Atoms object containing water molecules

    Returns
    -------
    np.ndarray
        Array of coordination numbers for oxygen atoms
    """
    atype = np.array(atoms.get_chemical_symbols())
    c_ids = np.where(atype == "O")[0]
    neigh_ids = np.where(atype == "H")[0]
    return calc_coord_number(atoms, c_ids, neigh_ids, 1.3)

calc_water_density(n, v)

Calculate water density using water molar mass (18.015 g/mol).

Parameters:

Name Type Description Default
n int or ndarray

Number of water molecules

 required
v float or ndarray

Volume in ų

 required

Returns:

Type Description
float or ndarray

Water density in g/cm³
Source code in toolbox/utils/utils.py 
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
def calc_water_density(
    n: Union[int, np.ndarray], v: Union[float, np.ndarray]
) -> Union[float, np.ndarray]:
    """
    Calculate water density using water molar mass (18.015 g/mol).

    Parameters
    ----------
    n : int or np.ndarray
        Number of water molecules
    v : float or np.ndarray
        Volume in ų

    Returns
    -------
    float or np.ndarray
        Water density in g/cm³
    """
    return calc_density(n, v, 18.015)

calc_water_number(rho, v)

Calculate number of water molecules from density and volume.

Parameters:

Name Type Description Default
rho float or ndarray

Water density in g/cm³

 required
v float or ndarray

Volume in ų

 required

Returns:

Type Description
int

Number of water molecules
Source code in toolbox/utils/utils.py 
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
def calc_water_number(rho: Union[float, np.ndarray], v: Union[float, np.ndarray]) -> int:
    """
    Calculate number of water molecules from density and volume.

    Parameters
    ----------
    rho : float or np.ndarray
        Water density in g/cm³
    v : float or np.ndarray
        Volume in ų

    Returns
    -------
    int
        Number of water molecules
    """
    return calc_number(rho, v, 18.015)

check_water(atoms)

Check if all water molecules have correct coordination (2 H per O).

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object containing water molecules

 required

Returns:

Type Description
bool

True if all water molecules are correct, False otherwise
Source code in toolbox/utils/utils.py 
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
def check_water(atoms: Atoms) -> bool:
    """
    Check if all water molecules have correct coordination (2 H per O).

    Parameters
    ----------
    atoms : Atoms
        ASE Atoms object containing water molecules

    Returns
    -------
    bool
        True if all water molecules are correct, False otherwise
    """
    cns = calc_water_coord_number(atoms)
    flags = cns == 2
    return False not in flags

dict_to_cp2k_input(input_dict)

Convert a dictionary to CP2K input file format.

Parameters:

Name Type Description Default
input_dict Dict[str, Any]

Dictionary containing CP2K input parameters

 required

Returns:

Type Description
str

Formatted CP2K input string
Source code in toolbox/utils/utils.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def dict_to_cp2k_input(input_dict: dict[str, Any]) -> str:
    """
    Convert a dictionary to CP2K input file format.

    Parameters
    ----------
    input_dict : Dict[str, Any]
        Dictionary containing CP2K input parameters

    Returns
    -------
    str
        Formatted CP2K input string
    """
    input_str = iterdict(input_dict, out_list=["\n"], loop_idx=0)
    s = "\n".join(input_str)
    s = s.strip("\n")
    return s

get_bins_from_bin_edge(bin_edges)

Get bin centers from bin edges.

Parameters:

Name Type Description Default
bin_edges ndarray or List[float]

Array of bin edges

 required

Returns:

Type Description
ndarray

Array of bin centers
Source code in toolbox/utils/utils.py 
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
def get_bins_from_bin_edge(bin_edges: Union[np.ndarray, List[float]]) -> np.ndarray:
    """
    Get bin centers from bin edges.

    Parameters
    ----------
    bin_edges : np.ndarray or List[float]
        Array of bin edges

    Returns
    -------
    np.ndarray
        Array of bin centers
    """
    bin_edges = np.reshape(bin_edges, (-1,))
    bins = bin_edges[:-1] + np.diff(bin_edges) / 2
    return bins

get_efields(delta_v, layer_thicknesses, eps)

Calculate electric fields from voltage differences and dielectric properties.

Parameters:

Name Type Description Default
delta_v float

Voltage difference

 required
layer_thicknesses List[float]

List of layer thicknesses

 required
eps List[float]

List of dielectric constants

 required

Returns:

Type Description
ndarray

Array of electric field values
Source code in toolbox/utils/utils.py 
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
def get_efields(delta_v: float, layer_thicknesses: list[float], eps: list[float]) -> np.ndarray:
    """
    Calculate electric fields from voltage differences and dielectric properties.

    Parameters
    ----------
    delta_v : float
        Voltage difference
    layer_thicknesses : List[float]
        List of layer thicknesses
    eps : List[float]
        List of dielectric constants

    Returns
    -------
    np.ndarray
        Array of electric field values
    """
    r_field = 1.0 / np.array(eps)
    _delta_v = np.sum(np.array(layer_thicknesses) * r_field)
    v_coeff = delta_v / _delta_v
    return r_field * v_coeff

iterdict(input_dict, out_list, loop_idx)

Recursively generate a list of strings for CP2K input file formatting.

This function processes a nested dictionary structure and converts it into a formatted list of strings suitable for generating CP2K input files.

Parameters:

Name Type Description Default
input_dict Dict[str, Any]

Dictionary containing CP2K input parameters

 required
out_list List[str]

List of strings for printing (modified in-place)

 required
loop_idx int

Record of loop levels in recursion

 required

Returns:

Type Description
List[str]

Formatted list of strings for CP2K input file
Source code in toolbox/utils/utils.py 
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
def iterdict(
    input_dict: dict[str, Any], out_list: list[str], loop_idx: int
) -> list[str]:
    """
    Recursively generate a list of strings for CP2K input file formatting.

    This function processes a nested dictionary structure and converts it into
    a formatted list of strings suitable for generating CP2K input files.

    Parameters
    ----------
    input_dict : Dict[str, Any]
        Dictionary containing CP2K input parameters
    out_list : List[str]
        List of strings for printing (modified in-place)
    loop_idx : int
        Record of loop levels in recursion

    Returns
    -------
    List[str]
        Formatted list of strings for CP2K input file
    """
    if len(out_list) == 0:
        out_list.append("\n")
    start_idx = len(out_list) - loop_idx - 2
    for k, v in input_dict.items():
        k = str(k)  # cast key into string
        # if value is dictionary
        if isinstance(v, dict):
            out_list.insert(-1 - loop_idx, "  " * loop_idx + "&" + k)
            out_list.insert(-1 - loop_idx, "  " * loop_idx + "&END " + k)
            iterdict(v, out_list, loop_idx + 1)
        # if value is list
        elif isinstance(v, list):
            if isinstance(v[0], dict):
                for _v in v:
                    out_list.insert(-1 - loop_idx, "  " * loop_idx + "&" + k)
                    out_list.insert(-1 - loop_idx, "  " * loop_idx + "&END " + k)
                    iterdict(_v, out_list, loop_idx + 1)
            else:
                for _v in v:
                    _v = str(_v)
                    out_list.insert(-1 - loop_idx, "  " * loop_idx + k + " " + _v)
        # if value is other type, e.g., int/float/str
        else:
            v = str(v)
            if k == "_":
                out_list[start_idx] = out_list[start_idx] + " " + v
            else:
                out_list.insert(-1 - loop_idx, "  " * loop_idx + k + " " + v)
    return out_list

load_dict(fname, fmt=None)

Load a dictionary from a file in various formats.

Parameters:

Name Type Description Default
fname str

Input filename

 required
fmt Optional[str]

File format (json, csv, pkl, hdf5, yaml). If None, inferred from extension.

 None

Returns:

Type Description
Dict[str, Any]

Loaded dictionary

Raises:

Type Description
KeyError

If the format is not supported
NotImplementedError

If HDF5 format is requested (not implemented)
Source code in toolbox/utils/utils.py 
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
def load_dict(fname: str, fmt: Optional[str] = None) -> dict[str, Any]:
    """
    Load a dictionary from a file in various formats.

    Parameters
    ----------
    fname : str
        Input filename
    fmt : Optional[str], optional
        File format (json, csv, pkl, hdf5, yaml). If None, inferred from extension.

    Returns
    -------
    Dict[str, Any]
        Loaded dictionary

    Raises
    ------
    KeyError
        If the format is not supported
    NotImplementedError
        If HDF5 format is requested (not implemented)
    """
    if fmt is None:
        fmt = os.path.splitext(fname)[1][1:]
    try:
        return globals()[f"load_dict_{fmt}"](fname)
    except KeyError as exc:
        raise KeyError(f"Unknown format {fmt}") from exc

load_dict_csv(fname)

Load dictionary from CSV file.

Parameters:

Name Type Description Default
fname str

Input filename

 required

Returns:

Type Description
Dict[str, str]

Loaded dictionary
Source code in toolbox/utils/utils.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
def load_dict_csv(fname: str) -> dict[str, str]:
    """
    Load dictionary from CSV file.

    Parameters
    ----------
    fname : str
        Input filename

    Returns
    -------
    Dict[str, str]
        Loaded dictionary
    """
    with open(fname, encoding="UTF-8") as f:
        data = csv.reader(f)
        d = {rows[0]: rows[1] for rows in data}
    return d

load_dict_hdf5(fname)

Load dictionary from HDF5 file (not implemented).

Parameters:

Name Type Description Default
fname str

Input filename

 required

Raises:

Type Description
NotImplementedError

This function is not implemented
Source code in toolbox/utils/utils.py 
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
def load_dict_hdf5(fname: str) -> dict[str, Any]:
    """
    Load dictionary from HDF5 file (not implemented).

    Parameters
    ----------
    fname : str
        Input filename

    Raises
    ------
    NotImplementedError
        This function is not implemented
    """
    raise NotImplementedError

load_dict_json(fname)

Load dictionary from JSON file.

Parameters:

Name Type Description Default
fname str

Input filename

 required

Returns:

Type Description
Dict[str, Any]

Loaded dictionary
Source code in toolbox/utils/utils.py 
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def load_dict_json(fname: str) -> dict[str, Any]:
    """
    Load dictionary from JSON file.

    Parameters
    ----------
    fname : str
        Input filename

    Returns
    -------
    Dict[str, Any]
        Loaded dictionary
    """
    with open(fname, encoding="UTF-8") as f:
        d = json.load(f)
    return d

load_dict_pkl(fname)

Load dictionary from pickle file.

Parameters:

Name Type Description Default
fname str

Input filename

 required

Returns:

Type Description
Dict[str, Any]

Loaded dictionary
Source code in toolbox/utils/utils.py 
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
def load_dict_pkl(fname: str) -> dict[str, Any]:
    """
    Load dictionary from pickle file.

    Parameters
    ----------
    fname : str
        Input filename

    Returns
    -------
    Dict[str, Any]
        Loaded dictionary
    """
    with open(fname, "rb") as f:
        d = pickle.load(f)
    return d

load_dict_yaml(fname)

Load dictionary from YAML file.

Parameters:

Name Type Description Default
fname str

Input filename

 required

Returns:

Type Description
Dict[str, Any]

Loaded dictionary
Source code in toolbox/utils/utils.py 
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
def load_dict_yaml(fname: str) -> dict[str, Any]:
    """
    Load dictionary from YAML file.

    Parameters
    ----------
    fname : str
        Input filename

    Returns
    -------
    Dict[str, Any]
        Loaded dictionary
    """
    with open(fname, encoding="UTF-8") as f:
        return yaml.safe_load(f)

safe_makedirs(dname)

Create directory if it doesn't exist.

Parameters:

Name Type Description Default
dname str

Directory path to create

 required
Source code in toolbox/utils/utils.py 
402
403
404
405
406
407
408
409
410
411
412
def safe_makedirs(dname: str) -> None:
    """
    Create directory if it doesn't exist.

    Parameters
    ----------
    dname : str
        Directory path to create
    """
    if not os.path.exists(dname):
        os.makedirs(dname)

Create symbolic link if it doesn't exist.

Parameters:

Name Type Description Default
src str

Source file path

 required
dst str

Destination path

 required
**kwargs

Additional arguments for os.symlink

 {}
Source code in toolbox/utils/utils.py 
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
def safe_symlink(src: str, dst: str, **kwargs) -> None:
    """
    Create symbolic link if it doesn't exist.

    Parameters
    ----------
    src : str
        Source file path
    dst : str
        Destination path
    **kwargs
        Additional arguments for os.symlink
    """
    with contextlib.suppress(OSError):
        os.symlink(src, dst, **kwargs)

save_dict(d, fname, fmt=None)

Save a dictionary to a file in various formats.

Args: d: Dictionary to save fname: Output filename fmt: File format (json, csv, pkl, hdf5, yaml). If None, inferred from extension.

Raises:

Type Description
KeyError: If the format is not supported
Source code in toolbox/utils/utils.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def save_dict(d: dict[str, Any], fname: str, fmt: Optional[str] = None) -> None:
    """
    Save a dictionary to a file in various formats.

    Args:
        d: Dictionary to save
        fname: Output filename
        fmt: File format (json, csv, pkl, hdf5, yaml). If None, inferred from extension.

    Raises
    ------
        KeyError: If the format is not supported
    """
    if fmt is None:
        fmt = os.path.splitext(fname)[1][1:]
    try:
        globals()[f"save_dict_{fmt}"](d, fname)
    except KeyError as exc:
        raise KeyError(f"Unknown format {fmt}") from exc

save_dict_csv(d, fname)

Save dictionary to CSV file.

Parameters:

Name Type Description Default
d Dict[str, Any]

Dictionary to save

 required
fname str

Output filename

 required
Source code in toolbox/utils/utils.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def save_dict_csv(d: dict[str, Any], fname: str) -> None:
    """
    Save dictionary to CSV file.

    Parameters
    ----------
    d : Dict[str, Any]
        Dictionary to save
    fname : str
        Output filename
    """
    with open(fname, "w", newline="", encoding="UTF-8") as f:
        writer = csv.DictWriter(f, fieldnames=d.keys())
        writer.writeheader()
        writer.writerow(d)

save_dict_hdf5(d, fname)

Save dictionary to HDF5 file.

Parameters:

Name Type Description Default
d Dict[str, Any]

Dictionary to save

 required
fname str

Output filename

 required
Source code in toolbox/utils/utils.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
def save_dict_hdf5(d: dict[str, Any], fname: str) -> None:
    """
    Save dictionary to HDF5 file.

    Parameters
    ----------
    d : Dict[str, Any]
        Dictionary to save
    fname : str
        Output filename
    """
    with h5py.File(fname, "a") as f:
        n = len(f)
        dts = f.create_group(f"{n:02d}")
        for k, v in d.items():
            dts.create_dataset(k, data=v)

save_dict_json(d, fname)

Save dictionary to JSON file.

Parameters:

Name Type Description Default
d Dict[str, Any]

Dictionary to save

 required
fname str

Output filename

 required
Source code in toolbox/utils/utils.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
def save_dict_json(d: dict[str, Any], fname: str) -> None:
    """
    Save dictionary to JSON file.

    Parameters
    ----------
    d : Dict[str, Any]
        Dictionary to save
    fname : str
        Output filename
    """
    with open(fname, "w", encoding="UTF-8") as f:
        json.dump(d, f, indent=4)

save_dict_pkl(d, fname)

Save dictionary to pickle file.

Parameters:

Name Type Description Default
d Dict[str, Any]

Dictionary to save

 required
fname str

Output filename

 required
Source code in toolbox/utils/utils.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
def save_dict_pkl(d: dict[str, Any], fname: str) -> None:
    """
    Save dictionary to pickle file.

    Parameters
    ----------
    d : Dict[str, Any]
        Dictionary to save
    fname : str
        Output filename
    """
    with open(fname, "wb") as f:
        pickle.dump(d, f)

save_dict_yaml(d, fname)

Save dictionary to YAML file.

Parameters:

Name Type Description Default
d Dict[str, Any]

Dictionary to save

 required
fname str

Output filename

 required
Source code in toolbox/utils/utils.py 
239
240
241
242
243
244
245
246
247
248
249
250
251
def save_dict_yaml(d: dict[str, Any], fname: str) -> None:
    """
    Save dictionary to YAML file.

    Parameters
    ----------
    d : Dict[str, Any]
        Dictionary to save
    fname : str
        Output filename
    """
    with open(fname, "w", encoding="UTF-8") as f:
        yaml.safe_dump(d, f)

Create a symbolic link.

Parameters:

Name Type Description Default
src str

Source file path

 required
_dst str

Destination path

 required
Source code in toolbox/utils/utils.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
def symlink(src: str, _dst: str) -> None:
    """
    Create a symbolic link.

    Parameters
    ----------
    src : str
        Source file path
    _dst : str
        Destination path
    """
    dst = os.path.abspath(_dst)
    os.symlink(src, dst)

update_dict(old_d, update_d)

Recursively update a dictionary with values from another dictionary.

Source: dpgen.generator.lib.cp2k

Parameters:

Name Type Description Default
old_d Dict[str, Any]

Original dictionary to be updated

 required
update_d Dict[str, Any]

Dictionary containing update values

 required
Source code in toolbox/utils/utils.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def update_dict(old_d: dict[str, Any], update_d: dict[str, Any]) -> None:
    """
    Recursively update a dictionary with values from another dictionary.

    Source: dpgen.generator.lib.cp2k

    Parameters
    ----------
    old_d : Dict[str, Any]
        Original dictionary to be updated
    update_d : Dict[str, Any]
        Dictionary containing update values
    """
    for k in update_d:
        if (
            k in old_d
            and isinstance(old_d[k], dict)
            and isinstance(update_d[k], collections.abc.Mapping)
        ):
            update_dict(old_d[k], update_d[k])
        else:
            old_d[k] = update_d[k]

wrap_water(atoms)

Make water molecules whole (keep O and H atoms together).

This function ensures that water molecules are not split across periodic boundaries. Works for pure water systems.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object containing water molecules

 required

Returns:

Type Description
Atoms

New Atoms object with wrapped water molecules
Source code in toolbox/utils/utils.py 
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
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
def wrap_water(atoms: Atoms) -> Atoms:
    """
    Make water molecules whole (keep O and H atoms together).

    This function ensures that water molecules are not split across periodic
    boundaries. Works for pure water systems.

    Parameters
    ----------
    atoms : Atoms
        ASE Atoms object containing water molecules

    Returns
    -------
    Atoms
        New Atoms object with wrapped water molecules
    """
    atoms = atoms.copy()
    oxygen_mask = atoms.symbols == "O"
    hydrogen_mask = atoms.symbols == "H"
    other_mask = np.logical_not(np.logical_or(oxygen_mask, hydrogen_mask))

    coords = wrap_positions(atoms.get_positions(), atoms.get_cell())
    cellpar = atoms.cell.cellpar()
    oxygen_coords = coords[oxygen_mask]
    hydrogen_coords = coords[hydrogen_mask]
    new_atoms = Atoms(cell=atoms.cell, pbc=atoms.pbc)
    for oxygen_coord in oxygen_coords:
        oxygen_coord = np.reshape(oxygen_coord, (1, 3))
        ds = distance_array(oxygen_coord, hydrogen_coords, box=cellpar)
        mask = ds.reshape(-1) < 1.3
        cn = np.sum(mask)
        coords_rel = hydrogen_coords[mask].reshape(-1, 3) - oxygen_coord
        coords_rel = minimize_vectors(coords_rel, box=cellpar)
        _coords = np.concatenate(
            (
                oxygen_coord,
                oxygen_coord + coords_rel,
            ),
            axis=0,
        )
        new_atoms.extend(Atoms(f"OH{cn}", positions=_coords))
        # remove selected hydrogen atoms
        hydrogen_coords = hydrogen_coords[~mask]
    assert len(hydrogen_coords) == 0
    new_atoms.extend(atoms[other_mask])
    return new_atoms

Machine Learning Module

TensorFlow Graph

toolbox.ml.tf_graph

TensorFlow graph module.

This module provides functionality for loading and visualizing TensorFlow graphs from DeepMD models.

Graph

from toolbox.ml import Graph.

Graph("graph.pb").run()

then you can access the graph with localhost:6006
Source code in toolbox/ml/tf_graph.py 
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
66
67
68
69
70
71
72
73
74
75
76
77
class Graph:
    """
    from toolbox.ml import Graph.

    Graph("graph.pb").run()
    # then you can access the graph with localhost:6006
    """

    def __init__(self, dp_model="graph.pb") -> None:
        """Initialize Graph.

        Parameters
        ----------
        dp_model : str, optional
            Path to Deep MD model file, by default "graph.pb"
        """
        self.model = dp_model

    def run(self, port=6006):
        """Run graph with TensorBoard visualization.

        Parameters
        ----------
        port : int, optional
            Port for TensorBoard server, by default 6006
        """
        graph = self._load_graph(self.model)
        with tf.Session(graph=graph) as sess:
            writer = tf.summary.FileWriter("dp_logs", sess.graph)
            writer.close()
        os.system(f"tensorboard --logdir=dp_logs --port={port:d}")

    @staticmethod
    def _load_graph(
        frozen_graph_filename, prefix: str = "load", default_tf_graph: bool = False
    ):
        """deepmd.infer.DeepEval._load_graph."""
        # We load the protobuf file from the disk and parse it to retrieve the
        # unserialized graph_def
        with tf.gfile.GFile(str(frozen_graph_filename), "rb") as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())

            if default_tf_graph:
                tf.import_graph_def(
                    graph_def,
                    input_map=None,
                    return_elements=None,
                    name=prefix,
                    producer_op_list=None,
                )
                graph = tf.get_default_graph()
            else:
                # Then, we can use again a convenient built-in function to import
                # a graph_def into the  current default Graph
                with tf.Graph().as_default() as graph:
                    tf.import_graph_def(
                        graph_def,
                        input_map=None,
                        return_elements=None,
                        name=prefix,
                        producer_op_list=None,
                    )

            return graph
__init__(dp_model='graph.pb')

Initialize Graph.

Parameters:

Name Type Description Default
dp_model str

Path to Deep MD model file, by default "graph.pb"

 'graph.pb'
Source code in toolbox/ml/tf_graph.py 
21
22
23
24
25
26
27
28
29
def __init__(self, dp_model="graph.pb") -> None:
    """Initialize Graph.

    Parameters
    ----------
    dp_model : str, optional
        Path to Deep MD model file, by default "graph.pb"
    """
    self.model = dp_model
run(port=6006)

Run graph with TensorBoard visualization.

Parameters:

Name Type Description Default
port int

Port for TensorBoard server, by default 6006

 6006
Source code in toolbox/ml/tf_graph.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
def run(self, port=6006):
    """Run graph with TensorBoard visualization.

    Parameters
    ----------
    port : int, optional
        Port for TensorBoard server, by default 6006
    """
    graph = self._load_graph(self.model)
    with tf.Session(graph=graph) as sess:
        writer = tf.summary.FileWriter("dp_logs", sess.graph)
        writer.close()
    os.system(f"tensorboard --logdir=dp_logs --port={port:d}")