annotate larch_select_paths.py @ 1:7fdca938d90c draft

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

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