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