Coverage for ase / io / nwchem / nwwriter.py: 78.80%
283 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3import os
4import random
5import string
6import warnings
7from copy import deepcopy
9import numpy as np
11from ase.calculators.calculator import KPoints, kpts2kpts
13_special_kws = ['center', 'autosym', 'autoz', 'theory', 'basis', 'xc', 'task',
14 'set', 'symmetry', 'label', 'geompar', 'basispar', 'kpts',
15 'bandpath', 'restart_kw', 'pretasks', 'charge']
17_system_type = {1: 'polymer', 2: 'surface', 3: 'crystal'}
20def _render_geom(atoms, params: dict) -> list[str]:
21 """Generate the geometry block
23 Parameters
24 ----------
25 atoms : Atoms
26 Geometry for the computation
27 params : dict
28 Parameter set for the computation
30 Returns
31 -------
32 geom : [str]
33 Geometry block to use in the computation
34 """
35 geom_header = ['geometry units angstrom']
36 for geomkw in ['center', 'autosym', 'autoz']:
37 geom_header.append(geomkw if params.get(geomkw) else 'no' + geomkw)
38 if 'geompar' in params:
39 geom_header.append(params['geompar'])
40 geom = [' '.join(geom_header)]
42 outpos = atoms.get_positions()
43 pbc = atoms.pbc
44 if np.any(pbc):
45 scpos = atoms.get_scaled_positions()
46 for i, pbci in enumerate(pbc):
47 if pbci:
48 outpos[:, i] = scpos[:, i]
49 npbc = pbc.sum()
50 cellpars = atoms.cell.cellpar()
51 geom.append(f' system {_system_type[npbc]} units angstrom')
52 if npbc == 3:
53 geom.append(' lattice_vectors')
54 for row in atoms.cell:
55 geom.append(' {:20.16e} {:20.16e} {:20.16e}'.format(*row))
56 else:
57 if pbc[0]:
58 geom.append(f' lat_a {cellpars[0]:20.16e}')
59 if pbc[1]:
60 geom.append(f' lat_b {cellpars[1]:20.16e}')
61 if pbc[2]:
62 geom.append(f' lat_c {cellpars[2]:20.16e}')
63 if pbc[1] and pbc[2]:
64 geom.append(f' alpha {cellpars[3]:20.16e}')
65 if pbc[0] and pbc[2]:
66 geom.append(f' beta {cellpars[4]:20.16e}')
67 if pbc[1] and pbc[0]:
68 geom.append(f' gamma {cellpars[5]:20.16e}')
69 geom.append(' end')
71 for i, atom in enumerate(atoms):
72 geom.append(' {:<2} {:20.16e} {:20.16e} {:20.16e}'
73 ''.format(atom.symbol, *outpos[i]))
74 symm = params.get('symmetry')
75 if symm is not None:
76 geom.append(f' symmetry {symm}')
77 geom.append('end')
78 return geom
81def _render_basis(theory, params: dict) -> list[str]:
82 """Infer the basis set block
84 Arguments
85 ---------
86 theory : str
87 Name of the theory used for the calculation
88 params : dict
89 Parameters for the computation
91 Returns
92 -------
93 [str]
94 List of input file lines for the basis block
95 """
97 # Break if no basis set is provided and non is applicable
98 if 'basis' not in params:
99 if theory in ['pspw', 'band', 'paw']:
100 return []
102 # Write the header section
103 if 'basispar' in params:
104 header = 'basis {} noprint'.format(params['basispar'])
105 else:
106 header = 'basis noprint'
107 basis_out = [header]
109 # Write the basis set for each atom type
110 basis_in = params.get('basis', '3-21G')
111 if isinstance(basis_in, str):
112 basis_out.append(f' * library {basis_in}')
113 else:
114 for symbol, ibasis in basis_in.items():
115 basis_out.append(f'{symbol:>4} library {ibasis}')
116 basis_out.append('end')
117 return basis_out
120_special_keypairs = [('nwpw', 'simulation_cell'),
121 ('nwpw', 'carr-parinello'),
122 ('nwpw', 'brillouin_zone'),
123 ('tddft', 'grad'),
124 ]
127def _render_brillouin_zone(array, name=None) -> list[str]:
128 out = [' brillouin_zone']
129 if name is not None:
130 out += [f' zone_name {name}']
131 template = ' kvector' + ' {:20.16e}' * array.shape[1]
132 for row in array:
133 out.append(template.format(*row))
134 out.append(' end')
135 return out
138def _render_bandpath(bp) -> list[str]:
139 if bp is None:
140 return []
141 out = ['nwpw']
142 out += _render_brillouin_zone(bp.kpts, name=bp.path)
143 out += [f' zone_structure_name {bp.path}',
144 'end',
145 'task band structure']
146 return out
149def _format_line(key, val) -> str:
150 if val is None:
151 return key
152 if isinstance(val, bool):
153 return f'{key} .{str(val).lower()}.'
154 else:
155 return ' '.join([key, str(val)])
158def _format_block(key, val, nindent=0) -> list[str]:
159 prefix = ' ' * nindent
160 prefix2 = ' ' * (nindent + 1)
161 if val is None:
162 return [prefix + key]
164 if not isinstance(val, dict):
165 return [prefix + _format_line(key, val)]
167 out = [prefix + key]
168 for subkey, subval in val.items():
169 if (key, subkey) in _special_keypairs:
170 if (key, subkey) == ('nwpw', 'brillouin_zone'):
171 out += _render_brillouin_zone(subval)
172 else:
173 out += _format_block(subkey, subval, nindent + 1)
174 else:
175 if isinstance(subval, dict):
176 subval = ' '.join([_format_line(a, b)
177 for a, b in subval.items()])
178 out.append(prefix2 + ' '.join([_format_line(subkey, subval)]))
179 out.append(prefix + 'end')
180 return out
183def _render_other(params) -> list[str]:
184 """Render other commands
186 Parameters
187 ----------
188 params : dict
189 Parameter set to be rendered
191 Returns
192 -------
193 out : [str]
194 Block defining other commands
195 """
196 out = []
197 for kw, block in params.items():
198 if kw in _special_kws:
199 continue
200 out += _format_block(kw, block)
201 return out
204def _render_set(set_params) -> list[str]:
205 """Render the commands for the set parameters
207 Parameters
208 ----------
209 set_params : dict
210 Parameters being set
212 Returns
213 -------
214 out : [str]
215 Block defining set commands
216 """
217 return ['set ' + _format_line(key, val) for key, val in set_params.items()]
220_gto_theories = ['tce', 'ccsd', 'tddft', 'scf', 'dft',
221 'direct_mp2', 'mp2', 'rimp2']
222_pw_theories = ['band', 'pspw', 'paw']
223_all_theories = _gto_theories + _pw_theories
226def _get_theory(params: dict) -> str:
227 """Infer the theory given the user-provided settings
229 Parameters
230 ----------
231 params : dict
232 Parameters for the computation
234 Returns
235 -------
236 theory : str
237 Theory directive to use
238 """
239 # Default: user-provided theory
240 theory = params.get('theory')
241 if theory is not None:
242 return theory
244 # Check if the user passed a theory to xc
245 xc = params.get('xc')
246 if xc in _all_theories:
247 return xc
249 # Check for input blocks that correspond to a particular level of
250 # theory. Correlated theories (e.g. CCSD) are checked first.
251 for kw in _gto_theories:
252 if kw in params:
253 return kw
255 # If the user passed an 'nwpw' block, then they want a plane-wave
256 # calculation, but what kind? If they request k-points, then
257 # they want 'band', otherwise assume 'pspw' (if the user wants
258 # to use 'paw', they will have to ask for it specifically).
259 nwpw = params.get('nwpw')
260 if nwpw is not None:
261 if 'monkhorst-pack' in nwpw or 'brillouin_zone' in nwpw:
262 return 'band'
263 return 'pspw'
265 # When all else fails, default to dft.
266 return 'dft'
269_xc_conv = dict(lda='slater pw91lda',
270 pbe='xpbe96 cpbe96',
271 revpbe='revpbe cpbe96',
272 rpbe='rpbe cpbe96',
273 pw91='xperdew91 perdew91',
274 )
277def _update_mult(magmom_tot: int, params: dict) -> None:
278 """Update parameters for multiplicity given the magnetic moment
280 For example, sets the number of open shells for SCF calculations
281 and the multiplicity for DFT calculations.
283 Parameters
284 ----------
285 magmom_tot : int
286 Total magnetic moment of the system
287 params : dict
288 Current set of parameters, will be modified
289 """
290 # Determine theory and multiplicity
291 theory = params['theory']
292 if magmom_tot == 0:
293 magmom_mult = 1
294 else:
295 magmom_mult = np.sign(magmom_tot) * (abs(magmom_tot) + 1)
297 # Adjust the kwargs for each type of theory
298 if 'scf' in params:
299 for kw in ['nopen', 'singlet', 'doublet', 'triplet', 'quartet',
300 'quintet', 'sextet', 'septet', 'octet']:
301 if kw in params['scf']:
302 break
303 else:
304 params['scf']['nopen'] = magmom_tot
305 elif theory in ['scf', 'mp2', 'direct_mp2', 'rimp2', 'ccsd', 'tce']:
306 params['scf'] = dict(nopen=magmom_tot)
308 if 'dft' in params:
309 if 'mult' not in params['dft']:
310 params['dft']['mult'] = magmom_mult
311 elif theory in ['dft', 'tddft']:
312 params['dft'] = dict(mult=magmom_mult)
314 if 'nwpw' in params:
315 if 'mult' not in params['nwpw']:
316 params['nwpw']['mult'] = magmom_mult
317 elif theory in ['pspw', 'band', 'paw']:
318 params['nwpw'] = dict(mult=magmom_mult)
321def _update_kpts(atoms, params) -> None:
322 """Converts top-level 'kpts' argument to native keywords
324 Parameters
325 ----------
326 atoms : Atoms
327 Input structure
328 params : dict
329 Current parameter set, will be updated
330 """
331 kpts = params.get('kpts')
332 if kpts is None:
333 return
335 nwpw = params.get('nwpw', {})
337 if 'monkhorst-pack' in nwpw or 'brillouin_zone' in nwpw:
338 raise ValueError("Redundant k-points specified!")
340 if isinstance(kpts, KPoints):
341 nwpw['brillouin_zone'] = kpts.kpts
342 elif isinstance(kpts, dict):
343 if kpts.get('gamma', False) or 'size' not in kpts:
344 nwpw['brillouin_zone'] = kpts2kpts(kpts, atoms).kpts
345 else:
346 nwpw['monkhorst-pack'] = ' '.join(map(str, kpts['size']))
347 elif isinstance(kpts, np.ndarray):
348 nwpw['brillouin_zone'] = kpts
349 else:
350 nwpw['monkhorst-pack'] = ' '.join(map(str, kpts))
352 params['nwpw'] = nwpw
355def _render_pretask(
356 this_step: dict,
357 previous_basis: list[str] | None,
358 wfc_path: str,
359 next_steps: list[dict],
360) -> tuple[list[str], list[str]]:
361 """Generate input file lines that perform a cheaper method first
363 Parameters
364 ----------
365 this_step: dict
366 Input parameters used to define the computation
367 previous_basis: [str], optional
368 Basis set block used in the previous step
369 wfc_path: str
370 Name of the wavefunction path
371 next_steps: [dict]
372 Parameters for the next steps in the calculation.
373 This function will adjust the next steps to read
374 and project from the wave functions written to disk by this step
375 if the basis set changes between this step and the next.
377 Returns
378 -------
379 output: [str]
380 Output lines for this task
381 this_basis: [str]
382 Basis set block used for this task
383 """
385 # Get the theory for the next step
386 next_step = next_steps[0]
387 next_theory = _get_theory(next_step)
388 next_step['theory'] = next_theory
389 out = []
390 if next_theory not in ['dft', 'mp2', 'direct_mp2', 'rimp2', 'scf']:
391 raise ValueError(f'Initial guesses not supported for {next_theory}')
393 # Determine the theory for this step
394 this_theory = _get_theory(this_step)
395 this_step['theory'] = this_theory
396 if this_theory not in ['dft', 'scf']:
397 raise ValueError('Initial guesses must use either dft or scf')
399 # Determine which basis set to use for this step. Our priorities
400 # 1. Basis defined explicitly in this step
401 # 2. Basis set of the previous step
402 # 3. Basis set of the target computation level
403 if 'basis' in this_step:
404 this_basis = _render_basis(this_theory, this_step)
405 elif previous_basis is not None:
406 this_basis = previous_basis.copy()
407 else:
408 # Use the basis for the last step
409 this_basis = _render_basis(next_theory, next_steps[-1])
411 # Determine the basis for the next step
412 # If not defined, it'll be the same as this step
413 if 'basis' in next_step:
414 next_basis = _render_basis(next_theory, next_step)
415 else:
416 next_basis = this_basis
418 # Set up projections if needed
419 if this_basis == next_basis:
420 out.append('\n'.join(this_basis))
421 proj_from = None # We do not need to project basis
422 else:
423 # Check for known limitations of NWChem
424 if this_theory != next_theory:
425 msg = 'Theories must be the same if basis are different. ' \
426 f'This step: {this_theory}//{this_basis} ' \
427 f'Next step: {next_theory}//{next_basis}'
428 if 'basis' not in this_step:
429 msg += f". Consider specifying basis in {this_step}"
430 raise ValueError(msg)
431 if not any('* library' in x for x in this_basis):
432 raise ValueError('We can only support projecting from systems '
433 'where all atoms share the same basis')
435 # Append a new name to this basis function by
436 # appending it as the first argument of the basis block
437 proj_from = f"smb_{len(next_steps)}"
438 this_basis[0] = f'basis {proj_from} {this_basis[0][6:]}'
439 out.append('\n'.join(this_basis))
441 # Point ao basis (the active basis set) to this new basis set
442 out.append(f'set "ao basis" {proj_from}')
444 # Insert a command to write wfcs from this computation to disk
445 if this_theory not in this_step:
446 this_step[this_theory] = {}
447 if 'vectors' not in this_step[this_theory]:
448 this_step[this_theory]['vectors'] = {}
449 this_step[this_theory]['vectors']['output'] = wfc_path
451 # Check if the initial theory changes
452 if this_theory != next_theory and \
453 'lindep:n_dep' not in this_step.get('set', {}):
454 warnings.warn('Loading initial guess may fail if you do not specify'
455 ' the number of linearly-dependent basis functions.'
456 ' Consider adding {"set": {"lindep:n_dep": 0}} '
457 f' to the step: {this_step}.')
459 # Add this to the input file along with a "task * ignore" command
460 out.extend(['\n'.join(_render_other(this_step)),
461 '\n'.join(_render_set(this_step.get('set', {}))),
462 f'task {this_theory} ignore'])
464 # Command to read the wavefunctions in the next step
465 # Theory used to get the wavefunctions may be different (mp2 uses SCF)
466 wfc_theory = 'scf' if 'mp2' in next_theory else next_theory
467 next_step = next_step
468 if wfc_theory not in next_step:
469 next_step[wfc_theory] = {}
470 if 'vectors' not in next_step[wfc_theory]:
471 next_step[wfc_theory]['vectors'] = {}
473 if proj_from is None:
474 # No need for projection
475 next_step[wfc_theory]['vectors']['input'] = wfc_path
476 else:
477 # Define that we should project from our basis set
478 next_step[wfc_theory]['vectors']['input'] \
479 = f'project {proj_from} {wfc_path}'
481 # Replace the name of the basis set to the default
482 out.append('set "ao basis" "ao basis"')
484 return out, this_basis
487def write_nwchem_in(fd, atoms, properties=None, echo=False, **params):
488 """Writes NWChem input file.
490 See :class:`~ase.calculators.nwchem.NWChem` for available params.
492 Parameters
493 ----------
494 fd
495 file descriptor
496 atoms
497 atomic configuration
498 properties
499 list of properties to compute; by default only the
500 calculation of the energy is requested
501 echo
502 if True include the `echo` keyword at the top of the file,
503 which causes the content of the input file to be included
504 in the output file
505 params
506 dict of instructions blocks to be included
507 """
508 # Copy so we can alter params w/o changing it for the function caller
509 params = deepcopy(params)
511 if properties is None:
512 properties = ['energy']
514 if 'stress' in properties:
515 if 'set' not in params:
516 params['set'] = {}
517 params['set']['includestress'] = True
519 task = params.get('task')
520 if task is None:
521 if 'stress' in properties or 'forces' in properties:
522 task = 'gradient'
523 else:
524 task = 'energy'
526 _update_kpts(atoms, params)
528 # Determine the theory for each step
529 # We determine the theory ahead of time because it is
530 # used when generating other parts of the input file (e.g., _get_mult)
531 theory = _get_theory(params)
532 params['theory'] = theory
534 for pretask in params.get('pretasks', []):
535 pretask['theory'] = _get_theory(pretask)
537 if 'xc' in params:
538 xc = _xc_conv.get(params['xc'].lower(), params['xc'])
539 if theory in ['dft', 'tddft']:
540 if 'dft' not in params:
541 params['dft'] = {}
542 params['dft']['xc'] = xc
543 elif theory in ['pspw', 'band', 'paw']:
544 if 'nwpw' not in params:
545 params['nwpw'] = {}
546 params['nwpw']['xc'] = xc
548 # Update the multiplicity given the charge of the system
549 magmom_tot = int(atoms.get_initial_magnetic_moments().sum())
550 _update_mult(magmom_tot, params)
551 for pretask in params.get('pretasks', []):
552 _update_mult(magmom_tot, pretask)
554 # Determine the header options
555 label = params.get('label', 'nwchem')
556 perm = os.path.abspath(params.pop('perm', label))
557 scratch = os.path.abspath(params.pop('scratch', label))
558 restart_kw = params.get('restart_kw', 'start')
559 if restart_kw not in ('start', 'restart'):
560 raise ValueError("Unrecognised restart keyword: {}!"
561 .format(restart_kw))
562 short_label = label.rsplit('/', 1)[-1]
563 if echo:
564 out = ['echo']
565 else:
566 out = []
568 # Defines the geometry and global options
569 out.extend([f'title "{short_label}"',
570 f'permanent_dir {perm}',
571 f'scratch_dir {scratch}',
572 f'{restart_kw} {short_label}',
573 '\n'.join(_render_set(params.get('set', {}))),
574 '\n'.join(_render_geom(atoms, params))])
576 # Add the charge if provided
577 if 'charge' in params:
578 out.append(f'charge {params["charge"]}')
580 # Define the memory separately from the other options so that it
581 # is defined before any initial wfc guesses are performed
582 memory_opts = params.pop('memory', None)
583 if memory_opts is not None:
584 out.extend(_render_other(dict(memory=memory_opts)))
586 # If desired, output commands to generate the initial wavefunctions
587 if 'pretasks' in params:
588 wfc_path = f'tmp-{"".join(random.choices(string.hexdigits, k=8))}.wfc'
589 pretasks = params['pretasks']
590 previous_basis = None
591 for this_ind, this_step in enumerate(pretasks):
592 new_out, previous_basis = _render_pretask(
593 this_step,
594 previous_basis,
595 wfc_path,
596 pretasks[this_ind + 1:] + [params]
597 )
598 out.extend(new_out)
600 # Finish output file with the commands to perform the desired computation
601 out.extend(['\n'.join(_render_basis(theory, params)),
602 '\n'.join(_render_other(params)),
603 f'task {theory} {task}',
604 '\n'.join(_render_bandpath(params.get('bandpath', None)))])
606 fd.write('\n\n'.join(out))