Remove code duplication. A bit new fragment-quality functional structure

Author: comcon1

src/fairmd/lipids/bin/evaluate_quality.py

@@ -60,44 +60,30 @@ def _evaluate_op_qualities(simulations) -> int:
         for lipidname, lipid in simulation.lipids.items():
             md_lipid_ops = simulation.op_data[lipidname]
 
+            # TODO: bb is merged into headgroup. But sn-s do not.. Cryptic rule.
+            # TODO: What to do with other types of lipids?
             fragments = mollib.get_fragments(lipid.mapping_dict)
-            fragment_qual_dict = {}
-            data_dict = {}
+            fragment_qual_perexp = {}
+            lipid_quality_perexp = {}
+            lipid_quality_perexp = {}
 
             for expid in simulation["EXPERIMENT"]["ORDERPARAMETER"].get(lipidname, []):
                 print(f"OP quality of simulation data in {simulation['path']}")
                 print(
                     f".. evaluating {lipidname} lipid using experimental data from {expid}",
                 )
-                OP_qual_data = {}
                 exp_lipid_ops = opexps.loc(expid).data[lipidname]
-                exp_error = 0.02  # TODO: hardcoded error value, should be taken from experiment data when available
-
-                for key, op_array_ in md_lipid_ops.items():
-                    OP_array = op_array_.copy()
-                    if key not in exp_lipid_ops:
-                        continue
-                    OP_exp_val = exp_lipid_ops[key][0]
-                    if not np.isnan(OP_exp_val):
-                        op_quality = qq.prob_op_within_trustinterval(
-                            op_exp=OP_exp_val,
-                            exp_error=exp_error,
-                            op_sim=OP_array[0],
-                            op_sim_sd=OP_array[2],
-                        )
-                        OP_array += [OP_exp_val, exp_error, op_quality]
-                    OP_qual_data[key] = OP_array
-
                 # save qualities of simulation-vs-experiment into a dictionary
-                data_dict[expid] = OP_qual_data
+                lipid_quality_perexp[expid] = qq.atomic_quality(exp_lipid_ops, md_lipid_ops)
 
                 # calculate quality for molecule fragments headgroup, sn-1, sn-2
-                # TODO: bb is merged into headgroup. But sn-s do not.. Cryptic rule.
-                # TODO: What to do with other types of lipids?
-                fragment_qual_dict[expid] = qq.fragment_quality(fragments, exp_lipid_ops, md_lipid_ops)
+                _frq = qq.atomic2fragment_quality(lipid_quality_perexp[expid], fragments)
+                _frw = qq.weights_of_fragments_in_data(fragments, exp_lipid_ops)
+                # dot product [ qualities * weights ]
+                fragment_qual_perexp[expid] = {k: _frq[k] * _frw[k] for k in fragments}
 
             # Experiment-merged fragment quality for the lipid
-            fragment_quality_merged = qq.fragment_quality_unite_multexp(lipidname, fragment_qual_dict, fragments)
+            fragment_quality_merged = qq.fragment_quality_unite_multexp(lipidname, fragment_qual_perexp, fragments)
             system_quality[lipidname] = fragment_quality_merged
 
             # Write FQ for the lipid
@@ -116,9 +102,9 @@ def _evaluate_op_qualities(simulations) -> int:
 
             # write into the OrderParameters_quality.json quality data file
             outfile1 = os.path.join(wdir, lipidname + "_OrderParameters_quality.json")
-            _round_quality_values(data_dict)
+            _round_quality_values(lipid_quality_perexp)
             with open(outfile1, "w") as f:
-                json.dump(data_dict, f, cls=CompactJSONEncoder)
+                json.dump(lipid_quality_perexp, f, cls=CompactJSONEncoder)
 
         system_qual_output = qq.systemQuality(system_quality, simulation)
         # make system quality file

src/fairmd/lipids/quality.py

@@ -4,7 +4,6 @@
 TODO: remove code duplication and commented code
 """
 
-import re
 import warnings
 
 import numpy as np
@@ -118,69 +117,32 @@ def atomic_quality(exp_op_data: dict, sim_op_data: dict):
     :return: dictionary of type {"nC nH": quality value}.
     """
     exp_error = 0.02  # TODO: hardcoded error value, should be taken from experiment data when available
+
+    # union of keys in exp_op_data and sim_op_data
+    all_keys = set(exp_op_data.keys()) | set(sim_op_data.keys())
     res_dict = {}
 
-    for key_exp, value_exp in exp_op_data.items():
-        if not np.isnan(value_exp[0]) and key_exp in sim_op_data:
-            q = prob_op_within_trustinterval(
-                op_exp=value_exp[0],
-                exp_error=exp_error,
-                op_sim=sim_op_data[key_exp][0],
-                op_sim_sd=sim_op_data[key_exp][2],
-            )
-            res_dict[key_exp] = q
+    for key in all_keys:
+        if key not in exp_op_data or key not in sim_op_data:
+            res_dict[key] = np.nan
+            continue
+        q = prob_op_within_trustinterval(
+            op_exp=exp_op_data[key][0],
+            exp_error=exp_error,
+            op_sim=sim_op_data[key][0],
+            op_sim_sd=sim_op_data[key][2],
+        )
+        res_dict[key] = q
 
     return res_dict
 
 
-def fragment_quality(fragments: dict, exp_op_data: dict, sim_op_data: dict):
-    """
-    Calculate quality for a fragmented molecule (times their weights in exp data).
-
-    Depends on the experiment file what fragments are in this dictionary.
-
-    :param fragments: dictionary of type {fragment:lists of unames}.
-    :param exp_op_data: dictionary of type {op_uname: [op_value]}.
-    :param sim_op_data: dictionary of type {op_uname: [op_value, op_sigma, op_sd]}.
-
-    :return: dictionary of type {fragment: quality value}.
-    """
-    fragment_weights = weights_of_fragments_in_data(fragments, exp_op_data)
-    exp_error = 0.02  # TODO: hardcoded error value, should be taken from experiment data when available
-
-    # empty dictionary with fragment names as keys
-    fragment_quality = dict.fromkeys(fragments.keys())
-
+def atomic2fragment_quality(atomic_qual_dict: dict, fragments: dict):
+    fqdict = dict.fromkeys(fragments.keys())
     for frg_name, frg_atoms in fragments.items():
-        E_sum = 0
-        AV_sum = 0
-        if fragment_weights[frg_name] == 0:
-            fragment_quality[frg_name] = np.nan
-            continue
-        for key_exp, value_exp in exp_op_data.items():
-            if (
-                key_exp.split()[0] in frg_atoms  # process for 1 fragm
-                and not np.isnan(value_exp[0])
-                and key_exp in sim_op_data
-                # If the last is not true, then simulation value is missing.
-                # This allows to happen for, e.g. CH3-groups in
-                # UA force fields as CH-bond cannot be reconstructed for this carbon.
-            ):
-                QE = prob_op_within_trustinterval(
-                    op_exp=value_exp[0],
-                    exp_error=exp_error,
-                    op_sim=sim_op_data[key_exp][0],
-                    op_sim_sd=sim_op_data[key_exp][2],
-                )
-                E_sum += QE
-                AV_sum += 1
-        if AV_sum > 0:
-            E_F = (E_sum / AV_sum) * fragment_weights[frg_name]
-            fragment_quality[frg_name] = E_F
-        else:
-            fragment_quality[frg_name] = np.nan
-
-    return fragment_quality
+        q_list = [v for k, v in atomic_qual_dict.items() if k.split(" ")[0] in frg_atoms]
+        fqdict[frg_name] = np.nanmean(q_list) if len(q_list) > 0 else np.nan
+    return fqdict
 
 
 def fragment_quality_unite_multexp(