def collect_final_total_magnitude(
scan_dir: Union[Path, str],
fixed_ns_path: Union[str, Path],
output_file: Union[Path, str] = None,
missing_ok: bool = True,
):
fixed_ns_path = make_path(fixed_ns_path).with_suffix(".fixed_ns.h5")
if output_file is None:
output_file = (Path("data") / "final_total_magnitude_{}".format(
fixed_ns_path.name.replace(".fixed_ns.h5", ""), ) /
(make_path(scan_dir).name + ".txt"))
def fetch(index, path, parameters):
_, data, _, _ = read_fixed_ns_total_magnitude_hdf5(path /
fixed_ns_path)
return data[-1]
return collect_values(
scan_dir,
[
fixed_ns_path,
],
output_file,
fetch,
missing_ok=missing_ok,
)
def collect_final_expval(
scan_dir: Union[Path, str],
expval: Union[Path, str],
output_file: Union[Path, str] = None,
node: int = 1,
dof: int = 1,
missing_ok: bool = True,
):
expval = make_path(expval)
if output_file is None:
folder_name = "expval_" + expval.name.rstrip(".exp.h5")
if not folder_name.startswith("final_"):
folder_name = "final_" + folder_name
output_file = Path("data") / (folder_name) / (
make_path(scan_dir).name + ".txt")
def fetch(index, path, parameters):
_, data = numpy.array(read_expval_hdf5(path / expval))
return data[-1].real, data[-1].imag
return collect_values(scan_dir, [expval],
output_file,
fetch,
missing_ok=missing_ok)
def create_slideshow(
files: List[Union[Path, str]],
output_file: Union[Path, str],
duration: float = 30.0,
):
if not len(files):
LOGGER.error("No image files specified! Nothing to do.")
return
files = [make_path(file).resolve() for file in files]
output_file = make_path(output_file).resolve()
with tempfile.TemporaryDirectory() as tmp:
with WorkingDir(tmp):
max_number_len = len(str(len(files)))
for i, file in enumerate(files):
os.symlink(
file,
("{:0" + str(max_number_len) + "d}{}").format(
i, file.suffix),
)
cmd = [
"ffmpeg",
"-r",
str(float(len(files)) / duration),
"-i",
"%0" + str(max_number_len) + "d" + files[0].suffix,
"output.mp4",
]
subprocess.run(cmd)
shutil.copy2("output.mp4", output_file)
def __init__(self, wave_function: Union[str, Path], operator: Union[str, Path]):
self.logger = get_logger(__name__ + ".ComputeExpectationValueStatic")
# compute required paths
self.operator = make_path(operator)
self.wave_function = make_path(wave_function)
self.expval = self.wave_function.with_name(
self.wave_function.stem + "_" + self.operator.stem,
).with_suffix(".exp.h5")
self.name = str(self.expval.with_suffix(""))
def direct_plot(
input_files: List[Union[str, Path]],
output_file_base: Union[str, Path],
plot_func: Callable[[], Tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]],
plotting_args: plot.PlotArgs2D = None,
extensions: List[str] = [".pdf", ".png"],
decorator_funcs: List[
Callable[[matplotlib.figure.Figure, matplotlib.axes.Axes], Any]
] = [],
):
if plotting_args is None:
plotting_args = plot.PlotArgs2D()
input_files = [make_path(input_file) for input_file in input_files]
output_file_base = make_path(output_file_base)
pickle_file = output_file_base.parent / ("." + output_file_base.stem + ".pickle")
def action_write_pickle(targets):
dirname = Path(targets[0]).parent
if not dirname.exists():
dirname.mkdir(parents=True)
with open(targets[0], "wb") as fptr:
pickle.dump([plotting_args, len(decorator_funcs)], fptr, protocol=3)
yield {
"name": f"direct_plot:{output_file_base.name}:pickle".replace("=", "_"),
"targets": [str(pickle_file)],
"clean": True,
"actions": [action_write_pickle],
}
def action_plot(targets):
fig, axes = plot_func()
for axis in axes:
plotting_args.apply(axis)
for decorator_func in decorator_funcs:
decorator_func(fig, axis)
for target in targets:
plot.save(fig, target)
plot.close_figure(fig)
yield {
"name": f"direct_plot:{output_file_base.name}:plot".replace("=", "_"),
"file_dep": [str(input_file) for input_file in input_files]
+ [str(pickle_file)],
"targets": [str(output_file_base) + ext for ext in extensions],
"clean": True,
"actions": [
action_plot,
],
}
def __init__(self, wave_function: Union[str, Path],
basis: Union[str, Path], **kwargs):
self.logger = get_logger(__name__ + ".NumberStateAnalysisStatic")
self.wave_function = make_path(wave_function)
self.basis = make_path(basis)
self.result = make_path(
kwargs.get(
"name",
self.wave_function.with_name(
self.wave_function.stem + "_" + self.basis.stem, ),
), ).with_suffix(".fixed_ns.h5")
self.name = str(self.result.with_suffix(""))
def __init__(
self,
name: str,
wave_function: Union[str, Path],
hamiltonian: Union[str, Path],
**kwargs,
):
self.name = name
self.wave_function = wave_function
self.hamiltonian = hamiltonian
self.logger = get_logger(__name__ + ".Propagate")
self.diag_gauge_oper: Optional[str] = kwargs.get(
"diag_gauge_oper", None)
self.flags, self.flag_list = create_flags(**kwargs)
self.flag_list += ["-rst", "restart", "-opr", "hamiltonian"]
if self.flags.get("gauge", "standard") == "diagonalization":
self.flag_list += ["-diag_gauge_oper", "gauge_oper"]
if self.flags["relax"]:
self.basename = "relaxation"
elif self.flags["improved_relax"]:
self.basename = "improved_relaxation"
elif self.flags["exact_diag"]:
self.basename = "exact_diagonalization"
else:
self.basename = "propagation"
# compute required paths
self.path_name = Path(self.name)
self.path_pickle = Path(self.name + ".prop_pickle")
self.path_wave_function = make_path(wave_function)
self.path_hamiltonian = make_path(hamiltonian)
self.path_diag_gauge_oper: Optional[Path] = (Path(
self.diag_gauge_oper) if self.diag_gauge_oper else None)
self.path_hdf5 = self.path_name / (
"exact_diag.h5" if self.flags["exact_diag"] else "propagate.h5")
if self.flags["exact_diag"]:
self.qdtk_files: List[str] = ["eigenenergies", "eigenvectors"]
else:
self.qdtk_files = ["final.wfn"]
if self.flags["psi"]:
self.qdtk_files.append("psi")
self.qdtk_files += kwargs.get("extra_files", [])
def collect_max_depletion(
scan_dir: Union[Path, str],
propagation_name: str = "propagate",
output_file: Union[Path, str] = None,
node: int = 1,
dof: int = 1,
missing_ok: bool = True,
):
if output_file is None:
output_file = (Path("data") / "max_depletion" /
(make_path(scan_dir).name + ".txt"))
def fetch(index, path, parameters):
_, data = read_natpop_hdf5(
path / propagation_name / "propagate.h5",
"natpop",
node=node,
dof=dof,
)
return (1 - data[:, 0]).max()
return collect_values(
scan_dir,
[
Path(propagation_name) / "propagate.h5",
],
output_file,
fetch,
missing_ok=missing_ok,
)
def scan_plot_energy(
scan_dir: Union[Path, str],
propagation: str = "propagate",
extensions: List[str] = [".png"],
**kwargs,
):
scan_dir = make_path(scan_dir)
plotting_args = PlotArgs2D.from_dict(kwargs)
selection = load_scan(scan_dir)
def plot_func(index, path, parameters):
del path
del parameters
time, _, energy, _ = read_output_hdf5(
scan_dir / "by_index" / str(index) / propagation / "propagate.h5",
"output",
)
fig, axis = plt.subplots(1, 1)
plot_energy(axis, time, energy)
return fig, [axis]
yield doit_plot_individual(
selection,
"energy",
[str(Path(propagation) / "propagate.h5")],
plot_func,
plotting_args,
extensions,
decorator_funcs=kwargs.get("decorator_funcs", []),
)
def collect_final_energy(
scan_dir: Union[Path, str],
propagation_name: str = "propagate",
output_file: Union[Path, str] = None,
missing_ok: bool = True,
):
if output_file is None:
output_file = (Path("data") / f"{propagation_name}_final_energy" /
(make_path(scan_dir).name + ".txt"))
def fetch(index, path, parameters):
_, _, energy, _ = read_output_hdf5(
path / propagation_name / "propagate.h5",
"output",
)
return energy[-1]
return collect_values(
scan_dir,
[
Path(propagation_name) / "propagate.h5",
],
output_file,
fetch,
missing_ok=missing_ok,
)
def __init__(
self,
psi: Union[str, Path],
index: int,
path: Union[Path, str] = None,
):
self.logger = get_logger(__name__ + ".FrameFromPsi")
self.psi = make_path(psi)
self.index = index
if path is None:
self.path = self.psi.stem + f"_{self.index}.wfn"
else:
self.path = make_path(path)
self.pickle_path = Path(str(self.path) + ".pickle")
def is_hdf5_path(path: Union[str, Path]) -> Tuple[bool, Path, str]:
"""Parse a path which potentially resides inside a HDF5 file
This function checks wether path points to a regular file on the disk. If
this is ``False``, the full path and an empty string are returned.
If path is a path to a HDF5 file any trailing components are considered to
denote the interior path inside the HDF5 file. In this case ``True``, the
path to the HDF5 file and the interior path are returned.
Args:
path (str): path to a file
Returns:
Tuple[bool, str, str]: whether this path is a path in a HDF5 file, path
to the file and interior path.
"""
path = make_path(path).resolve()
built_path = Path(path.parts[0])
for part in path.parts[1:]:
built_path /= part
if is_hdf5_file(built_path):
return True, built_path, str(path.relative_to(built_path))
return False, path, ""
def read_fixed_ns_ascii(
path: Union[str, Path],
) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
path = make_path(path)
with open(path) as fptr:
num_coefficients = (len(fptr.readline().split()) - 1) // 3
names = (["time"] +
["real_" + str(i) for i in range(num_coefficients)] +
["imag_" + str(i) for i in range(num_coefficients)])
usecols = [i for i in range(2 * num_coefficients + 1)]
data = pandas.read_csv(
path,
delim_whitespace=True,
header=None,
names=names,
usecols=usecols,
)[names].values
times, indices = numpy.unique(data[:, 0], return_index=True)
return (
times,
data[indices, 1:num_coefficients + 1],
data[indices, num_coefficients + 1:],
)
def read_dmat_evecs_spf_ascii(
path: Union[str, Path],
) -> Tuple[numpy.ndarray, numpy.ndarray]:
path = make_path(path)
with open(path) as fp:
m = (len(fp.readline().strip().split()) - 2) // 2
names = ["time", "index"]
names_real = []
names_imag = []
for i in range(m):
names.append("real_" + str(i))
names.append("imag_" + str(i))
names_real.append("real_" + str(i))
names_imag.append("imag_" + str(i))
df = pandas.read_csv(path, delim_whitespace=True, header=None, names=names)
evecs = df[names_real].values + 1j * df[names_imag].values
num_indices = 0
t_0 = df["time"][0]
for t in df["time"]:
if t_0 != t:
break
num_indices += 1
times = df["time"].values[::num_indices]
return times, evecs.T.reshape(m, len(times), num_indices)
def create_slideshow(
images: List[Union[Path, str]],
output_file: Union[Path, str],
duration: float = 20.0,
):
output_file = make_path(output_file)
pickle_file = output_file.with_suffix(output_file.suffix + ".pickle")
def action_write_pickle(duration, targets):
make_path(targets[0]).parent.mkdir(parents=True, exist_ok=True)
obj = {"duration": duration}
with open(targets[0], "wb") as fptr:
pickle.dump(obj, fptr, protocol=3)
yield {
"name": f"slideshow:{str(pickle_file)}:pickle".replace("=", "_"),
"targets": [
pickle_file,
],
"clean": True,
"actions": [(action_write_pickle, [duration])],
}
def action_create_slideshow(duration, images, targets):
video.create_slideshow(images, targets[0], duration)
yield {
"name": f"slideshow:{str(output_file)}:create".replace("=", "_"),
"targets": [
output_file,
],
"clean": True,
"actions": [(action_create_slideshow, [duration, images])],
}
def load_scan(path: Union[str, Path]) -> ParameterSelection:
"""Load all parameter sets of a parameter scan.
Args:
path: Path to the parameter scan containing the file ``scan.pickle``.
"""
path = make_path(path)
with open(path / "scan.pickle", "rb") as fptr:
obj = pickle.load(fptr)
return ParameterSelection((parameter for parameter in obj), path)
def read_dmat_evals_ascii(
path: Union[str, Path],
) -> Tuple[numpy.ndarray, numpy.ndarray]:
path = make_path(path)
with open(path) as fp:
m = len(fp.readline().strip().split()) - 1
names = ["time"] + ["orbital_" + str(i) for i in range(m)]
df = pandas.read_csv(path, delim_whitespace=True, header=None, names=names)
return df["time"].values, df[names[1:]].values
def __init__(
self,
parameters: Iterable[Parameters],
path: Union[str, Path] = None,
indices: Iterable[int] = None,
):
if indices is None:
self.parameters = [(i, p) for i, p in enumerate(parameters)]
else:
self.parameters = list(zip(indices, parameters))
self.path = None if path is None else make_path(path).resolve()
def __init__(
self,
scan_dir: Union[Path, str],
propagation: str = "propagate",
node: int = 1,
dof: int = 1,
missing_ok: bool = True,
output_file: Union[Path, str] = None,
):
self.scan_dir = make_path(scan_dir)
self.propagation = propagation
self.node = node
self.dof = dof
self.missing_ok = missing_ok
if output_file is None:
self.output_file = (Path("data") / f"natpop_{node}_{dof}" /
(self.scan_dir.name.replace("=", "_") + ".h5"))
else:
self.output_file = make_path(output_file)
def __init__(
self,
name: str,
simulations: List[Simulation],
working_dir: Union[str, Path] = None,
):
self.name = name
self.simulations = simulations
self.working_dir = (Path(name).resolve()
if working_dir is None else make_path(working_dir))
self.logger = get_logger(__name__ + ".SimulationSet")
def __init__(self, psi: Union[str, Path], operator: Union[str, Path], **kwargs):
self.logger = get_logger(__name__ + ".ComputeExpectationValue")
self.unique_name = kwargs.get("unique_name", False)
# compute required paths
self.operator = make_path(operator)
self.psi = make_path(psi)
if self.unique_name:
self.expval = self.psi.with_name(
self.psi.stem + "_" + self.operator.stem,
).with_suffix(".exp.h5")
else:
self.expval = (self.psi.parent / self.operator.stem).with_suffix(".exp.h5")
if "wave_function" in kwargs:
self.wave_function = make_path(kwargs["wave_function"])
else:
self.wave_function = make_path(
kwargs.get("wave_function", self.psi.parent / "final"),
).with_suffix(".wfn")
self.name = str(self.expval.with_suffix(""))
def scan_plot_expval(
scan_dir: Union[Path, str],
expval: Union[Path, str],
extensions: List[str] = [
".png",
],
**kwargs,
):
scan_dir = make_path(scan_dir)
expval = make_path(expval)
kwargs["coefficient"] = kwargs.get("coefficient", 1.0)
plotting_args = PlotArgs2D.from_dict(kwargs)
selection = load_scan(scan_dir)
def plot_func(index, path, parameters):
del path
del parameters
data = read_expval_hdf5(
str((scan_dir / "by_index" / str(index) /
expval).with_suffix(".exp.h5")), )
fig, axis = plt.subplots(1, 1)
plot_expval(axis, *data, **kwargs)
return fig, [axis]
return doit_plot_individual(
selection,
f"expval_{str(expval)}".replace("/", "_"),
[str(expval.with_suffix(".exp.h5"))],
plot_func,
plotting_args,
extensions,
decorator_funcs=kwargs.get("decorator_funcs", []),
extra_args={"coefficient": kwargs["coefficient"]},
)
def read_g2_hdf5(
path: Union[str, Path],
interior_path: str = "/g2",
) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
path = make_path(path)
with h5py.File(path, "r") as fptr:
return (
fptr[interior_path]["time"][:],
fptr[interior_path]["x1"][:],
fptr[interior_path]["x2"][:],
fptr[interior_path]["real"][:] +
1j * fptr[interior_path]["imag"][:],
)
def save_pdf(
figure: matplotlib.figure.Figure,
path: Union[str, Path],
crop: bool = True,
optimize: bool = True,
):
path = make_path(path)
figure.savefig(path)
if crop:
crop_pdf(path)
if optimize:
optimize_pdf(path)
def save(
figure: matplotlib.figure.Figure,
path: Union[str, Path],
crop: bool = False,
optimize: bool = False,
):
path = make_path(path)
if path.suffix == ".pdf":
save_pdf(figure, path, crop, optimize)
return
if crop:
figure.savefig(path, bbox_inches="tight")
else:
figure.savefig(path)
def scan_gpop_slideshow(
scan_dir: Union[Path, str],
dof: int = 1,
duration: float = 20.0,
):
scan_dir = make_path(scan_dir)
yield create_slideshow(
list_files(
scan_dir / "plots" / ("gpop_" + str(dof)),
[
".png",
],
),
(Path("videos") / ("gpop_" + str(dof)) / scan_dir.name).with_suffix(
scan_dir.suffix + ".mp4", ),
duration,
)
def read_dmat_spfrep_ascii(
path: Union[str, Path],
) -> Tuple[numpy.ndarray, numpy.ndarray]:
path = make_path(path)
df = pandas.read_csv(
str(path),
delim_whitespace=True,
header=None,
names=["time", "i", "j", "real", "imag"],
)
time = numpy.unique(df["time"].values)
num_times = len(time)
num_i = len(numpy.unique(df["i"].values))
num_j = len(numpy.unique(df["j"].values))
elements = df["real"].values + 1j * df["imag"].values
return time, elements.reshape((num_times, num_i, num_j))
def scan_natpop_slideshow(
scan_dir: Union[Path, str],
node: int = 1,
dof: int = 1,
duration: float = 20.0,
):
scan_dir = make_path(scan_dir)
yield create_slideshow(
list_files(
scan_dir / "plots" / f"natpop_{node}_{dof}",
[
".png",
],
),
(Path("videos") / (f"natpop_{node}_{dof}") /
scan_dir.name).with_suffix(scan_dir.suffix + ".mp4", ),
duration,
)
def __init__(
self,
name: str,
parameters: Parameters,
db_path: Union[str, Path],
variable_name: str = "db",
compute: bool = True,
):
db_path = make_path(db_path)
if not db_path.is_absolute():
self.db_path = Path(sys.argv[0]).resolve().parent / db_path
else:
self.db_path = db_path
self.name = name
self.parameters = parameters
self.variable_name = variable_name
self.compute = compute
self.path = Path(self.name)
def scan_plot_entropy(
scan_dir: Union[Path, str],
propagation: str = "propagate",
node: int = 1,
dof: int = 1,
extensions: List[str] = [
".png",
],
**kwargs,
):
scan_dir = make_path(scan_dir)
plotting_args = PlotArgs2D.from_dict(kwargs)
selection = load_scan(scan_dir)
def plot_func(index, path, parameters):
del path
del parameters
data = read_natpop(
str(scan_dir / "by_index" / str(index) / propagation /
"propagate.h5") + "/natpop",
node=node,
dof=dof,
)
entropy = compute_entropy(data[1])
fig, axis = plt.subplots(1, 1)
plot_entropy(axis, data[0], entropy)
return fig, [axis]
return doit_plot_individual(
selection,
"entropy",
[str(Path(propagation) / "propagate.h5")],
plot_func,
plotting_args,
extensions,
decorator_funcs=kwargs.get("decorator_funcs", []),
)
