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