Skip to content

Commit

Permalink
pep8 compliance
Browse files Browse the repository at this point in the history
  • Loading branch information
elinscott committed Oct 10, 2024
1 parent 690df56 commit bc4bc44
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 63 deletions.
31 changes: 17 additions & 14 deletions docs/_static/tutorials/fetch.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,14 @@ def extract(filename_in, filename_out, start=0, end=None, heading=None):

flines = flines[start:end]

# Find the shortest leading whitespace and strip this from all lines (as otherwise the markdown will be rendered as a code block)
# Find the shortest leading whitespace and strip this from all lines (as otherwise the markdown will be rendered
# as a code block)
min_indent = min(len(line) - len(line.lstrip()) for line in flines if line.strip())
flines = [line if line == "\n" else line[min_indent:]for line in flines]

# Manual conversion of double spaces to /
flines = [line[:-2] + '\ \n' if (line.endswith(' \n')
and not line.strip().startswith('-')) else line for line in flines]
# Manual conversion of double spaces to "\"
flines = [line[:-2] + '\\ \n' if (line.endswith(' \n')
and not line.strip().startswith('-')) else line for line in flines]

with open(filename_out, 'w') as fd:
fd.writelines(flines)
Expand All @@ -48,15 +49,17 @@ def extract(filename_in, filename_out, start=0, end=None, heading=None):

# Tutorial 2
extract('../../../tutorials/tutorial_2/si_wannierize.md', 'tutorial_2/md_excerpts/si_wannierize.md', start=18)
extract('../../../tutorials/tutorial_2/si_ki.md', 'tutorial_2/md_excerpts/si_ki_wannierize.md', heading="Wannierize")
extract('../../../tutorials/tutorial_2/si_ki.md', 'tutorial_2/md_excerpts/si_ki_fold.md', heading="Fold To Supercell")
extract('../../../tutorials/tutorial_2/si_ki.md', 'tutorial_2/md_excerpts/si_ki_wannierize.md',
heading="Wannierize")
extract('../../../tutorials/tutorial_2/si_ki.md', 'tutorial_2/md_excerpts/si_ki_fold.md',
heading="Fold To Supercell")
extract('../../../tutorials/tutorial_2/si_ki.md', 'tutorial_2/md_excerpts/si_ki_screening.md', 51, 61)
extract('../../../tutorials/tutorial_2/si_ki.md',
'tutorial_2/md_excerpts/si_ki_postproc.md', heading="Unfold And Interpolate")
extract('../../../tutorials/tutorial_2/si_ki.md', 'tutorial_2/md_excerpts/si_ki_postproc.md',
heading="Unfold And Interpolate")

# Tutorial 3
extract('../../../tutorials/tutorial_3/01-ki/zno.md',
'tutorial_3/md_excerpts/zno_wannierize_section.md', heading='Wannierize')
extract('../../../tutorials/tutorial_3/01-ki/zno.md', 'tutorial_3/md_excerpts/zno_wannierize_section.md',
heading='Wannierize')
extract('../../../tutorials/tutorial_3/01-ki/zno.md', 'tutorial_3/md_excerpts/zno_w2kc.md', -4, -3)
extract('../../../tutorials/tutorial_3/01-ki/zno.md', 'tutorial_3/md_excerpts/zno_ham.md', -3, -2)
copy('../../../tutorials/tutorial_3/01-ki/01-koopmans-dfpt/Koopmans_DFPT_bandstructure.png', 'tutorial_3/')
Expand All @@ -66,13 +69,13 @@ def extract(filename_in, filename_out, start=0, end=None, heading=None):

# Tutorial 5
extract('../../../tutorials/tutorial_5/01-train/h2o_train.md', 'tutorial_5/md_excerpts/train.md', 45, 92)
extract('../../../tutorials/tutorial_5/02-predict/h2o_predict.md',
'tutorial_5/md_excerpts/predict.md', heading='Calculate Screening Via DSCF')
extract('../../../tutorials/tutorial_5/02-predict/h2o_predict.md', 'tutorial_5/md_excerpts/predict.md',
heading='Calculate Screening Via DSCF')
extract('../../../tutorials/tutorial_5/03-test/h2o_test.md', 'tutorial_5/md_excerpts/test.md', 45, 96)

# Tutorial 6
extract('../../../tutorials/tutorial_6/cri3.md', 'tutorial_6/md_excerpts/cri3_wannierize_spin_up.md', 20, 35)
extract('../../../tutorials/tutorial_6/cri3.md', 'tutorial_6/md_excerpts/cri3_wannierize_spin_down.md', 43, 55)
extract('../../../tutorials/tutorial_6/cri3.md', 'tutorial_6/md_excerpts/cri3_final.md', heading='02-ki_final')
extract('../../../tutorials/tutorial_6/cri3.md',
'tutorial_6/md_excerpts/cri3_postproc.md', heading='Unfold And Interpolate')
extract('../../../tutorials/tutorial_6/cri3.md', 'tutorial_6/md_excerpts/cri3_postproc.md',
heading='Unfold And Interpolate')
27 changes: 15 additions & 12 deletions src/koopmans/workflows/_wannierize.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ def __init__(self, *args, force_nspin2=False, scf_kgrid=None, **kwargs):
if not spin:
num_bands_occ /= 2
divs = self.projections.divisions(spin)
cumulative_divs = [sum(divs[:i+1]) for i in range(len(divs))]
cumulative_divs = [sum(divs[:i + 1]) for i in range(len(divs))]
if num_bands_occ not in cumulative_divs:
message = 'The provided `Wannier90` projections are not commensurate with the number of ' \
'electrons; divide your list of projections into sublists that represent blocks ' \
Expand All @@ -104,8 +104,8 @@ def __init__(self, *args, force_nspin2=False, scf_kgrid=None, **kwargs):
pass

else:
raise NotImplementedError('`WannierizeWorkflow` only supports setting `init_orbitals` and `init_empty_orbitals` '
'to `mlwfs`/`projwfs`/`kohn-sham`')
raise NotImplementedError('`WannierizeWorkflow` only supports setting `init_orbitals` and '
'`init_empty_orbitals` to `mlwfs`/`projwfs`/`kohn-sham`')

# Spin-polarization
self._force_nspin2 = force_nspin2
Expand Down Expand Up @@ -163,7 +163,8 @@ def _run(self):
for block in self.projections:
wannierize_block_subworkflow = WannierizeBlockWorkflow.fromparent(
self, force_nspin2=self._force_nspin2, block=block)
wannierize_block_subworkflow.name = f'Wannierize {block.name.replace("_", " ").replace("block", "Block")}'
wannierize_block_subworkflow.name = \
f'Wannierize {block.name.replace("_", " ").replace("block", "Block")}'
wannierize_block_subworkflow.run()

# Store the results
Expand Down Expand Up @@ -211,10 +212,11 @@ def _run(self):
u_matrices_files[label] = FilePointer(merge_u_proc, merge_u_proc.outputs.dst_file)

# Merging the wannier centers files
merge_centers_proc = MergeProcess(merge_function=partial(merge_wannier_centers_file_contents, atoms=self.atoms),
src_files=[(calc, Path(calc.prefix + '_centres.xyz'))
for calc in src_calcs],
dst_file=prefix + f'{emp_label}_centres.xyz')
merge_centers_proc = MergeProcess(
merge_function=partial(merge_wannier_centers_file_contents, atoms=self.atoms),
src_files=[(calc, Path(calc.prefix + '_centres.xyz'))
for calc in src_calcs],
dst_file=prefix + f'{emp_label}_centres.xyz')
merge_centers_proc.name = f'merge_{label}_wannier_centers'
self.run_process(merge_centers_proc)
centers_files[label] = FilePointer(merge_centers_proc, merge_centers_proc.outputs.dst_file)
Expand Down Expand Up @@ -259,7 +261,8 @@ def _run(self):

# Link the save directory so that the bands calculation can use the old density
self.link(calc_nscf, (calc_nscf.parameters.outdir / calc_nscf.parameters.prefix).with_suffix('.save'),
calc_pw_bands, (calc_pw_bands.parameters.outdir / calc_pw_bands.parameters.prefix).with_suffix('.save'))
calc_pw_bands,
(calc_pw_bands.parameters.outdir / calc_pw_bands.parameters.prefix).with_suffix('.save'))
self.run_calculator(calc_pw_bands)

# Calculate a projected DOS
Expand Down Expand Up @@ -292,9 +295,9 @@ def _run(self):

# Prepare the band structures for plotting
ax = None
labels = ['explicit'] \
+ [f'interpolation ({c.absolute_directory.parent.name.split("-", 2)[-1].replace("block-", "block ").replace("spin-", "spin ").replace("-",", ")})'
for c in selected_calcs]
labels = ['explicit'] + [f'interpolation ({c.absolute_directory.parent.name})' for c in selected_calcs]
labels = [l.split("-", 2)[-1].replace("block-", "block ").replace("spin-", "spin ").replace("-", ", ")
for l in labels]
color_cycle = plt.rcParams['axes.prop_cycle']()
bsplot_kwargs_list = []
colors = {}
Expand Down
Loading

0 comments on commit bc4bc44

Please sign in to comment.