Mercurial > repos > muon-spectroscopy-computational-project > larch_plot

comparison larch_plot.py @ 0:886949a03377 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/MaterialsGalaxy/larch-tools/tree/main/larch_plot commit 5be486890442dedfb327289d597e1c8110240735
author muon-spectroscopy-computational-project
date Tue, 14 Nov 2023 15:35:36 +0000
parents
children 002c18a3e642
comparison
equal deleted inserted replaced
-1:000000000000 0:886949a03377
1 import json
2 import sys
3
4 from common import read_groups
5
6 import matplotlib
7 import matplotlib.pyplot as plt
8
9 import numpy as np
10
11
12 Y_LABELS = {
13 "norm": r"x$\mu$(E), normalised",
14 "dmude": r"d(x$\mu$(E))/dE, normalised",
15 "chir_mag": r"|$\chi$(r)|",
16 }
17
18
19 def main(dat_files: "list[str]", plot_settings: "list[dict]"):
20 groups = list(read_groups(dat_files))
21
22 for i, settings in enumerate(plot_settings):
23 data_list = []
24 e0_min = None
25 e0_max = None
26 variable = settings["variable"]["variable"]
27 x_min = settings["variable"]["energy_min"]
28 x_max = settings["variable"]["energy_max"]
29 plot_path = f"plots/{i}_{variable}.png"
30 plt.figure()
31
32 for group in groups:
33 label = group.athena_params.annotation or group.athena_params.id
34 if variable == "chir_mag":
35 x = group.r
36 energy_format = None
37 else:
38 x = group.energy
39 energy_format = settings["variable"]["energy_format"]
40 if energy_format == "relative":
41 e0 = group.athena_params.bkg.e0
42 e0_min = find_relative_limit(e0_min, e0, min)
43 e0_max = find_relative_limit(e0_max, e0, max)
44
45 y = getattr(group, variable)
46 if x_min is None and x_max is None:
47 plt.plot(x, y, label=label)
48 else:
49 data_list.append({"x": x, "y": y, "label": label})
50
51 if variable != "chir_mag" and energy_format == "relative":
52 if x_min is not None:
53 x_min += e0_min
54 if x_max is not None:
55 x_max += e0_max
56
57 if x_min is not None or x_max is not None:
58 for data in data_list:
59 index_min = None
60 index_max = None
61 x = data["x"]
62 if x_min is not None:
63 index_min = max(np.searchsorted(x, x_min) - 1, 0)
64 if x_max is not None:
65 index_max = min(np.searchsorted(x, x_max) + 1, len(x))
66 plt.plot(
67 x[index_min:index_max],
68 data["y"][index_min:index_max],
69 label=data["label"],
70 )
71
72 plt.xlim(x_min, x_max)
73
74 save_plot(variable, plot_path)
75
76
77 def find_relative_limit(e0_min: "float|None", e0: float, function: callable):
78 if e0_min is None:
79 e0_min = e0
80 else:
81 e0_min = function(e0_min, e0)
82 return e0_min
83
84
85 def save_plot(y_type: str, plot_path: str):
86 plt.grid(color="r", linestyle=":", linewidth=1)
87 plt.xlabel("Energy (eV)")
88 plt.ylabel(Y_LABELS[y_type])
89 plt.legend()
90 plt.savefig(plot_path, format="png")
91 plt.close("all")
92
93
94 if __name__ == "__main__":
95 # larch imports set this to an interactive backend, so need to change it
96 matplotlib.use("Agg")
97
98 dat_files = sys.argv[1]
99 input_values = json.load(open(sys.argv[2], "r", encoding="utf-8"))
100
101 main(dat_files.split(","), input_values["plots"])