annotate larch_select_paths.py @ 0:2e827836f0ad draft

planemo upload for repository https://github.com/MaterialsGalaxy/larch-tools/tree/main/larch_select_paths commit 5be486890442dedfb327289d597e1c8110240735

1 import csv
2 import json
3 import os
4 import re
5 import sys
6 from zipfile import ZIP_DEFLATED, ZipFile
7 
8 
9 class GDSWriter:
10     def __init__(self, default_variables: "dict[str, dict]"):
11         self.default_properties = {
12             "s02": {"name": "s02"},
13             "e0": {"name": "e0"},
14             "deltar": {"name": "alpha*reff"},
15             "sigma2": {"name": "sigma2"},
16         }
17         self.rows = [
18             f"{'id':>4s}, {'name':>24s}, {'value':>5s}, {'expr':>4s}, "
19             f"{'vary':>4s}\n"
20         ]
21         self.names = set()
22 
23         for property in self.default_properties:
24             name = self.default_properties[property]["name"]
25             value = default_variables[property]["value"]
26             vary = default_variables[property]["vary"]
27             is_common = default_variables[property]["is_common"]
28 
29             self.default_properties[property]["value"] = value
30             self.default_properties[property]["vary"] = vary
31             self.default_properties[property]["is_common"] = is_common
32 
33             if is_common:
34                 self.append_gds(name=name, value=value, vary=vary)
35 
36     def append_gds(
37         self,
38         name: str,
39         value: float = 0.,
40         expr: str = None,
41         vary: bool = True,
42         label: str = "",
43     ):
44         """Append a single GDS variable to the list of rows, later to be
45         written to file.
46 
47         Args:
48             name (str): Name of the GDS variable.
49             value (float, optional): Starting value for variable.
50                 Defaults to 0.
51             expr (str, optional): Expression for setting the variable.
52                 Defaults to None.
53             vary (bool, optional): Whether the variable is optimised during the
54                 fit. Defaults to True.
55             label (str, optional): Label to keep variables for different FEFF
56                 directories distinct. Defaults to "".
57         """
58         formatted_name = name if (label is None) else label + name
59         formatted_name = formatted_name.replace("*reff", "")
60         if not expr:
61             expr = "    "
62 
63         if formatted_name in self.names:
64             raise ValueError(f"{formatted_name} already used as variable name")
65         self.names.add(formatted_name)
66 
67         self.rows.append(
68             f"{len(self.rows):4d}, {formatted_name:>24s}, {str(value):>5s}, "
69             f"{expr:>4s}, {str(vary):>4s}\n"
70         )
71 
72     def parse_gds(
73         self,
74         property_name: str,
75         variable_name: str = None,
76         path_variable: dict = None,
77         directory_label: str = None,
78         path_label: str = None,
79     ) -> str:
80         """Parse and append a row defining a GDS variable for a particular
81         path.
82 
83         Args:
84             property_name (str): The property to which the variable
85                 corresponds. Should be a key in `self.default_properties`.
86             variable_name (str, optional): Custom name for this variable.
87                 Defaults to None.
88             path_variable (dict, optional): Dictionary defining the GDS
89                 settings for this path's variable. Defaults to None.
90             directory_label (str, optional): Label to indicate paths from a
91                 separate directory. Defaults to None.
92             path_label (str, optional): Label indicating the atoms involved in
93                 this path. Defaults to None.
94 
95         Returns:
96             str: Either `variable_name`, the name used as a default globally
97                 for this `property_name`, or an automatically generated unique
98                 name.
99         """
00         if variable_name:
01             self.append_gds(
02                 name=variable_name,
03                 value=path_variable["value"],
04                 expr=path_variable["expr"],
05                 vary=path_variable["vary"],
06             )
07             return variable_name
08         elif self.default_properties[property_name]["is_common"]:
09             return self.default_properties[property_name]["name"]
10         else:
11             auto_name = self.default_properties[property_name]["name"]
12             if directory_label:
13                 auto_name += f"_{directory_label}"
14             if path_label:
15                 auto_name += f"_{path_label.lower().replace('.', '')}"
16 
17             self.append_gds(
18                 name=auto_name,
19                 value=self.default_properties[property_name]["value"],
20                 vary=self.default_properties[property_name]["vary"],
21             )
22             return auto_name
23 
24     def write(self):
25         """Write GDS rows to file.
26         """
27         with open("gds.csv", "w") as out:
28             out.writelines(self.rows)
29 
30 
31 class PathsWriter:
32     def __init__(self, default_variables: "dict[str, dict]"):
33         self.rows = [
34             f"{'id':>4s}, {'filename':>24s}, {'label':>24s}, {'s02':>3s}, "
35             f"{'e0':>4s}, {'sigma2':>24s}, {'deltar':>10s}\n"
36         ]
37         self.gds_writer = GDSWriter(default_variables=default_variables)
38 
39     def parse_feff_output(
40         self,
41         paths_file: str,
42         selection: "dict[str, str|list]",
43         directory_label: str = "",
44     ):
45         """Parse selected paths from CSV summary and define GDS variables.
46 
47         Args:
48             paths_file (str): CSV summary filename.
49             selection (dict[str, str|list]): Dictionary indicating which paths
50                 to select, and how to define their variables.
51             directory_label (str, optional): Label to indicate paths from a
52                 separate directory. Defaults to "".
53         """
54         paths = selection["paths"]
55         path_values_ids = [path_value["id"] for path_value in paths]
56 
57         with open(paths_file) as file:
58             reader = csv.reader(file)
59             for row in reader:
60                 id_match = re.search(r"\d+", row[0])
61                 if id_match:
62                     path_id = int(id_match.group())
63                     filename = row[0].strip()
64                     path_label = row[-2].strip()
65                     variables = {}
66 
67                     if path_id in path_values_ids:
68                         path_value = paths[path_values_ids.index(path_id)]
69                         for property in self.gds_writer.default_properties:
70                             variables[property] = self.gds_writer.parse_gds(
71                                 property_name=property,
72                                 variable_name=path_value[property]["name"],
73                                 path_variable=path_value[property],
74                                 directory_label=directory_label,
75                                 path_label=path_label,
76                             )
77                         self.parse_selected_path(
78                             filename=filename,
79                             path_label=path_label,
80                             directory_label=directory_label,
81                             **variables,
82                         )
83                     elif selection["selection"] == "all" or int(row[-1]):
84                         path_value = None
85                         for property in self.gds_writer.default_properties:
86                             variables[property] = self.gds_writer.parse_gds(
87                                 property_name=property,
88                                 directory_label=directory_label,
89                                 path_label=path_label,
90                             )
91                         self.parse_selected_path(
92                             filename=filename,
93                             path_label=path_label,
94                             directory_label=directory_label,
95                             **variables,
96                         )
97 
98     def parse_selected_path(
99         self,
00         filename: str,
01         path_label: str,
02         directory_label: str = "",
03         s02: str = "s02",
04         e0: str = "e0",
05         sigma2: str = "sigma2",
06         deltar: str = "alpha*reff",
07     ):
08         """Format and append row representing a selected FEFF path.
09 
10         Args:
11             filename (str): Name of the underlying FEFF path file, without
12                 parent directory.
13             path_label (str): Label indicating the atoms involved in this path.
14             directory_label (str, optional): Label to indicate paths from a
15                 separate directory. Defaults to "".
16             s02 (str, optional): Electron screening factor variable name.
17                 Defaults to "s02".
18             e0 (str, optional): Energy shift variable name. Defaults to "e0".
19             sigma2 (str, optional): Mean squared displacement variable name.
20                 Defaults to "sigma2".
21             deltar (str, optional): Change in path length variable.
22                 Defaults to "alpha*reff".
23         """
24         if directory_label:
25             filename = os.path.join(directory_label, filename)
26             label = f"{directory_label}.{path_label}"
27         else:
28             filename = os.path.join("feff", filename)
29             label = path_label
30 
31         self.rows.append(
32             f"{len(self.rows):>4d}, {filename:>24s}, {label:>24s}, "
33             f"{s02:>3s}, {e0:>4s}, {sigma2:>24s}, {deltar:>10s}\n"
34         )
35 
36     def write(self):
37         """Write selected path and GDS rows to file.
38         """
39         self.gds_writer.write()
40         with open("sp.csv", "w") as out:
41             out.writelines(self.rows)
42 
43 
44 def main(input_values: dict):
45     """Select paths and define GDS parameters.
46 
47     Args:
48         input_values (dict): All input values from the Galaxy tool UI.
49 
50     Raises:
51         ValueError: If a FEFF label is not unique.
52     """
53     default_variables = input_values["variables"]
54 
55     writer = PathsWriter(default_variables=default_variables)
56 
57     if len(input_values["feff_outputs"]) == 1:
58         feff_output = input_values["feff_outputs"][0]
59         writer.parse_feff_output(
60             paths_file=feff_output["paths_file"],
61             selection=feff_output["selection"],
62         )
63     else:
64         zfill_length = len(str(len(input_values["feff_outputs"])))
65         labels = set()
66         with ZipFile("merged.zip", "x", ZIP_DEFLATED) as zipfile_out:
67             for i, feff_output in enumerate(input_values["feff_outputs"]):
68                 label = feff_output.pop("label") or str(i + 1).zfill(
69                     zfill_length
70                 )
71                 if label in labels:
72                     raise ValueError(f"Label '{label}' is not unique")
73                 labels.add(label)
74 
75                 writer.parse_feff_output(
76                     directory_label=label,
77                     paths_file=feff_output["paths_file"],
78                     selection=feff_output["selection"],
79                 )
80 
81                 with ZipFile(feff_output["paths_zip"]) as z:
82                     for zipinfo in z.infolist():
83                         if zipinfo.filename != "feff/":
84                             zipinfo.filename = zipinfo.filename[5:]
85                             z.extract(member=zipinfo, path=label)
86                             zipfile_out.write(
87                                 os.path.join(label, zipinfo.filename)
88                             )
89 
90     writer.write()
91 
92 
93 if __name__ == "__main__":
94     input_values = json.load(open(sys.argv[1], "r", encoding="utf-8"))
95     main(input_values)