Skip to content

Commit

Permalink
ruff format example [skip ci]
Browse files Browse the repository at this point in the history
  • Loading branch information
jcmgray committed Jan 31, 2025
1 parent 2b001e7 commit b6259ba
Showing 1 changed file with 18 additions and 8 deletions.
26 changes: 18 additions & 8 deletions docs/examples/ex_mpi_expm_evo.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
And will display slightly different output for each that explains the
three modes.
"""

import quimb as qu
from mpi4py import MPI

Expand All @@ -24,8 +25,9 @@

# this makes the function print some information when called
# - in order to be pickled is has to be located in the main package
ham_heis_verbose = qu.utils.Verbosify(qu.ham_heis,
highlight='ownership', mpi=True)
ham_heis_verbose = qu.utils.Verbosify(
qu.ham_heis, highlight="ownership", mpi=True
)

H = qu.Lazy(ham_heis_verbose, n=n, sparse=True, shape=shape)

Expand All @@ -36,7 +38,7 @@
# evolve the system, processes split 'hard' work (slepc computations)
# - should see each worker gets given a different ownership rows
# - but all end up with the results.
evo = qu.Evolution(psi0, H, method='expm', expm_backend='slepc')
evo = qu.Evolution(psi0, H, method="expm", expm_backend="slepc")
evo.update_to(5)


Expand All @@ -47,14 +49,22 @@

dims = [2] * n
bsz = 5
logneg_subsys_verbose = qu.utils.Verbosify(qu.logneg_subsys,
highlight='sysb', mpi=True)
logneg_subsys_verbose = qu.utils.Verbosify(
qu.logneg_subsys, highlight="sysb", mpi=True
)

# each process only computes its own fraction of these
# - should see each process calls logneg with different ``sysb``.
fs = [pool.submit(logneg_subsys_verbose, evo.pt, dims=dims,
sysa=range(0, bsz), sysb=range(i, i + bsz))
for i in range(bsz, n - bsz)]
fs = [
pool.submit(
logneg_subsys_verbose,
evo.pt,
dims=dims,
sysa=range(0, bsz),
sysb=range(i, i + bsz),
)
for i in range(bsz, n - bsz)
]

# but then the results are comminucated to everyone
rs = [f.result() for f in fs]
Expand Down

0 comments on commit b6259ba

Please sign in to comment.