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

changeset 1:002c18a3e642 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 1cf6d7160497ba58fe16a51f00d088a20934eba6
author muon-spectroscopy-computational-project
date Wed, 06 Dec 2023 13:04:06 +0000
parents 886949a03377
children 59d0d15a40ef
files common.py larch_plot.py larch_plot.xml macros.xml
diffstat 4 files changed, 181 insertions(+), 94 deletions(-) [+]
line wrap: on
line diff
--- a/common.py	Tue Nov 14 15:35:36 2023 +0000
+++ b/common.py	Wed Dec 06 13:04:06 2023 +0000
@@ -7,38 +7,145 @@
 
 
 def get_group(athena_group: AthenaGroup, key: str = None) -> Group:
+    group_keys = list(athena_group._athena_groups.keys())
     if key is None:
-        group_keys = list(athena_group._athena_groups.keys())
         key = group_keys[0]
-    return extract_athenagroup(athena_group._athena_groups[key])
+    else:
+        key = key.replace("-", "_")
+
+    try:
+        return extract_athenagroup(athena_group._athena_groups[key])
+    except KeyError as e:
+        raise KeyError(f"{key} not in {group_keys}") from e
+
+
+def read_all_groups(dat_file: str, key: str = None) -> "dict[str, Group]":
+    # Cannot rely on do_ABC as _larch is None
+    athena_group = read_athena(
+        dat_file,
+        do_preedge=False,
+        do_bkg=False,
+        do_fft=False,
+    )
+    all_groups = {}
+    for key in athena_group._athena_groups.keys():
+        group = get_group(athena_group, key)
+        pre_edge_with_defaults(group=group)
+        xftf_with_defaults(group=group)
+        all_groups[key] = group
+
+    return all_groups
+
+
+def read_group(dat_file: str, key: str = None):
+    # Cannot rely on do_ABC as _larch is None
+    athena_group = read_athena(
+        dat_file,
+        do_preedge=False,
+        do_bkg=False,
+        do_fft=False,
+    )
+    group = get_group(athena_group, key)
+    pre_edge_with_defaults(group=group)
+    xftf_with_defaults(group=group)
+    return group
 
 
-def read_group(dat_file: str, key: str = None, xftf_params: dict = None):
-    athena_group = read_athena(dat_file)
-    group = get_group(athena_group, key)
-    bkg_parameters = group.athena_params.bkg
-    print(group.athena_params.fft)
-    print(group.athena_params.fft.__dict__)
-    pre_edge(
-        group,
-        e0=bkg_parameters.e0,
-        pre1=bkg_parameters.pre1,
-        pre2=bkg_parameters.pre2,
-        norm1=bkg_parameters.nor1,
-        norm2=bkg_parameters.nor2,
-        nnorm=bkg_parameters.nnorm,
-        make_flat=bkg_parameters.flatten,
+def pre_edge_with_defaults(group: Group, settings: dict = None):
+    merged_settings = {}
+    try:
+        bkg_parameters = group.athena_params.bkg
+    except AttributeError as e:
+        print(f"Cannot load group.athena_params.bkg from group:\n{e}")
+        bkg_parameters = None
+
+    keys = (
+        ("e0", "e0", None),
+        ("pre1", "pre1", None),
+        ("pre2", "pre2", None),
+        ("norm1", "nor1", None),
+        ("norm2", "nor2", None),
+        ("nnorm", "nnorm", None),
+        ("make_flat", "flatten", None),
+        ("step", "step", None),
+        # This cannot be read from file as it is not stored by Larch (0.9.71)
+        # ("nvict", "nvict", None),
     )
-    autobk(group)
-    if xftf_params is None:
-        xftf(group)
-    else:
-        print(xftf_params)
-        xftf(group, **xftf_params)
-        xftf_details = Group()
-        setattr(xftf_details, "call_args", xftf_params)
-        group.xftf_details = xftf_details
-    return group
+    for key, parameters_key, default in keys:
+        extract_attribute(
+            merged_settings, key, bkg_parameters, parameters_key, default
+        )
+
+    if settings:
+        for k, v in settings.items():
+            if k == "nvict":
+                print(
+                    "WARNING: `nvict` can be used for pre-edge but is not "
+                    "saved to file, so value used will not be accessible in "
+                    "future operations using this Athena .prj"
+                )
+            merged_settings[k] = v
+
+    print(f"Pre-edge normalization with {merged_settings}")
+    try:
+        pre_edge(group, **merged_settings)
+    except Warning as e:
+        raise Warning(
+            "Unable to perform pre-edge fitting with:\n\n"
+            f"energy:\n{group.energy}\n\nmu:{group.mu}\n\n"
+            "Consider checking the correct columns have been extracted"
+        ) from e
+    autobk(group, pre_edge_kws=merged_settings)
+
+
+def xftf_with_defaults(group: Group, settings: dict = None):
+    merged_settings = {}
+    try:
+        fft_parameters = group.athena_params.fft
+    except AttributeError as e:
+        print(f"Cannot load group.athena_params.fft from group:\n{e}")
+        fft_parameters = None
+
+    keys = (
+        ("kmin", "kmin", 0),
+        ("kmax", "kmax", 20),
+        ("dk", "dk", 1),
+        ("kweight", "kw", 2),
+        ("kweight", "kweight", 2),
+        ("window", "kwindow", "kaiser"),
+    )
+    for key, parameters_key, default in keys:
+        extract_attribute(
+            merged_settings, key, fft_parameters, parameters_key, default
+        )
+
+    if settings:
+        for k, v in settings.items():
+            merged_settings[k] = v
+
+    print(f"XFTF with {merged_settings}")
+    xftf(group, **merged_settings)
+    xftf_details = Group()
+    setattr(xftf_details, "call_args", merged_settings)
+    group.xftf_details = xftf_details
+
+
+def extract_attribute(
+    merged_settings: dict,
+    key: str,
+    parameters_group: Group,
+    parameters_key: str,
+    default: "str|int" = None,
+):
+    if parameters_group is not None:
+        try:
+            merged_settings[key] = getattr(parameters_group, parameters_key)
+            return
+        except AttributeError:
+            pass
+
+    if default is not None:
+        merged_settings[key] = default
 
 
 def read_groups(dat_files: "list[str]", key: str = None) -> Iterable[Group]:
--- a/larch_plot.py	Tue Nov 14 15:35:36 2023 +0000
+++ b/larch_plot.py	Wed Dec 06 13:04:06 2023 +0000
@@ -9,10 +9,12 @@
 import numpy as np
 
 
-Y_LABELS = {
+AXIS_LABELS = {
     "norm": r"x$\mu$(E), normalised",
     "dmude": r"d(x$\mu$(E))/dE, normalised",
     "chir_mag": r"|$\chi$(r)|",
+    "energy": "Energy (eV)",
+    "distance": "r (ang)",
 }
 
 
@@ -21,39 +23,30 @@
 
     for i, settings in enumerate(plot_settings):
         data_list = []
-        e0_min = None
-        e0_max = None
-        variable = settings["variable"]["variable"]
-        x_min = settings["variable"]["energy_min"]
-        x_max = settings["variable"]["energy_max"]
-        plot_path = f"plots/{i}_{variable}.png"
+        x_variable = "energy"
+        y_variable = settings["variable"]["variable"]
+        x_min = settings["variable"]["x_limit_min"]
+        x_max = settings["variable"]["x_limit_max"]
+        y_min = settings["variable"]["y_limit_min"]
+        y_max = settings["variable"]["y_limit_max"]
+        plot_path = f"plots/{i}_{y_variable}.png"
         plt.figure()
 
         for group in groups:
-            label = group.athena_params.annotation or group.athena_params.id
-            if variable == "chir_mag":
+            params = group.athena_params
+            label = params.annotation or params.file or params.id
+            if y_variable == "chir_mag":
+                x_variable = "distance"
                 x = group.r
-                energy_format = None
             else:
                 x = group.energy
-                energy_format = settings["variable"]["energy_format"]
-                if energy_format == "relative":
-                    e0 = group.athena_params.bkg.e0
-                    e0_min = find_relative_limit(e0_min, e0, min)
-                    e0_max = find_relative_limit(e0_max, e0, max)
 
-            y = getattr(group, variable)
+            y = getattr(group, y_variable)
             if x_min is None and x_max is None:
                 plt.plot(x, y, label=label)
             else:
                 data_list.append({"x": x, "y": y, "label": label})
 
-        if variable != "chir_mag" and energy_format == "relative":
-            if x_min is not None:
-                x_min += e0_min
-            if x_max is not None:
-                x_max += e0_max
-
         if x_min is not None or x_max is not None:
             for data in data_list:
                 index_min = None
@@ -70,22 +63,15 @@
                 )
 
         plt.xlim(x_min, x_max)
+        plt.ylim(y_min, y_max)
 
-        save_plot(variable, plot_path)
+        save_plot(x_variable, y_variable, plot_path)
 
 
-def find_relative_limit(e0_min: "float|None", e0: float, function: callable):
-    if e0_min is None:
-        e0_min = e0
-    else:
-        e0_min = function(e0_min, e0)
-    return e0_min
-
-
-def save_plot(y_type: str, plot_path: str):
+def save_plot(x_type: str, y_type: str, plot_path: str):
     plt.grid(color="r", linestyle=":", linewidth=1)
-    plt.xlabel("Energy (eV)")
-    plt.ylabel(Y_LABELS[y_type])
+    plt.xlabel(AXIS_LABELS[x_type])
+    plt.ylabel(AXIS_LABELS[y_type])
     plt.legend()
     plt.savefig(plot_path, format="png")
     plt.close("all")
--- a/larch_plot.xml	Tue Nov 14 15:35:36 2023 +0000
+++ b/larch_plot.xml	Wed Dec 06 13:04:06 2023 +0000
@@ -4,7 +4,7 @@
         <!-- version of underlying tool (PEP 440) -->
         <token name="@TOOL_VERSION@">0.9.71</token>
         <!-- version of this tool wrapper (integer) -->
-        <token name="@WRAPPER_VERSION@">0</token>
+        <token name="@WRAPPER_VERSION@">1</token>
         <!-- citation should be updated with every underlying tool version -->
         <!-- typical fields to update are version, month, year, and doi -->
         <token name="@TOOL_CITATION@">10.1088/1742-6596/430/1/012007</token>
@@ -38,14 +38,20 @@
                     <option value="chir_mag">Magnitude of χ(r)</option>
                 </param>
                 <when value="norm">
-                    <expand macro="energy_limits"/>
+                    <expand macro="plot_limits_energy"/>
+                    <param name="y_limit_min" type="float" label="Minimum xμ(E)" optional="true" help="If set, plot will be limited to this value on the y axis."/>
+                    <param name="y_limit_max" type="float" label="Maximum xμ(E)" optional="true" help="If set, plot will be limited to this value on the y axis."/>
                 </when>
                 <when value="dmude">
-                    <expand macro="energy_limits"/>
+                    <expand macro="plot_limits_energy"/>
+                    <param name="y_limit_min" type="float" label="Minimum d(xμ(E))/dE" optional="true" help="If set, plot will be limited to this value on the y axis."/>
+                    <param name="y_limit_max" type="float" label="Maximum d(xμ(E))/dE" optional="true" help="If set, plot will be limited to this value on the y axis."/>
                 </when>
                 <when value="chir_mag">
-                    <param name="energy_min" type="float" label="Minimum r (ang)" optional="true" help="If set, data will be cropped below this value in angstroms."/>
-                    <param name="energy_max" type="float" label="Maximum r (ang)" optional="true" help="If set, data will be cropped above this value in angstroms."/>
+                    <param name="x_limit_min" type="float" label="Minimum r (ang)" optional="true" help="If set, plot will be limited to this value on the x axis."/>
+                    <param name="x_limit_max" type="float" label="Maximum r (ang)" optional="true" help="If set, plot will be limited to this value on the x axis."/>
+                    <param name="y_limit_min" type="float" label="Minimum |χ(r)|" optional="true" help="If set, plot will be limited to this value on the y axis."/>
+                    <param name="y_limit_max" type="float" label="Maximum |χ(r)|" optional="true" help="If set, plot will be limited to this value on the y axis."/>
                 </when>
             </conditional>
         </repeat>
@@ -56,15 +62,23 @@
         </collection>
     </outputs>
     <tests>
+        <!-- 1: plot types -->
         <test expect_num_outputs="1">
             <param name="dat_files" value="test.prj"/>
             <param name="variable" value="norm"/>
-            <param name="energy_format" value="absolute"/>
-            <param name="energy_min" value="7000"/>
             <param name="variable" value="dmude"/>
-            <param name="energy_format" value="relative"/>
-            <param name="energy_max" value="10"/>
-            <output_collection name="plot_collection" type="list" count="2"/>
+            <param name="variable" value="chir_mag"/>
+            <output_collection name="plot_collection" type="list" count="3"/>
+        </test>
+        <!-- 2: plot limits -->
+        <test expect_num_outputs="1">
+            <param name="dat_files" value="test.prj"/>
+            <param name="variable" value="norm"/>
+            <param name="x_limit_min" value="7000"/>
+            <param name="x_limit_max" value="7200"/>
+            <param name="y_limit_min" value="0"/>
+            <param name="y_limit_max" value="1"/>
+            <output_collection name="plot_collection" type="list" count="1"/>
         </test>
     </tests>
     <help><![CDATA[
--- a/macros.xml	Tue Nov 14 15:35:36 2023 +0000
+++ b/macros.xml	Wed Dec 06 13:04:06 2023 +0000
@@ -1,26 +1,6 @@
 <macros>
-    <xml name="energy_limits">
-        <param name="energy_format" type="select" display="radio" label="Energy limits" help="Whether to limit the energy relative to the absorption edge or with absolute values.">
-            <option value="relative" selected="true">Relative</option>
-            <option value="absolute">Absolute</option>
-        </param>
-        <param name="energy_min" type="float" label="Minimum energy (eV)" optional="true" help="If set, data will be cropped below this value in electron volts."/>
-        <param name="energy_max" type="float" label="Maximum energy (eV)" optional="true" help="If set, data will be cropped above this value in electron volts."/>
-    </xml>
-    <xml name="xftf_params">
-        <param argument="kmin" type="float" value="0" min="0.0" help="Minimum k value."/>
-        <param argument="kmax" type="float" value="20" min="0.0" help="Maximum k value."/>
-        <param argument="kweight" type="float" value="2" help="Exponent for weighting spectra by raising k to this power."/>
-        <param argument="dk" type="float" value="4" help="Tapering parameter for Fourier Transform window."/>
-        <param argument="window" type="select" help="Fourier Transform window type.">
-            <option value="hanning">Hanning (cosine-squared taper)</option>
-            <option value="parzen">Parzen (linear taper)</option>
-            <option value="welch">Welch (quadratic taper)</option>
-            <option value="gaussian">Gaussian function window</option>
-            <option value="sine">Sine function window</option>
-            <option value="kaiser" selected="true">Kaiser-Bessel function-derived window</option>
-        </param>
-        <param argument="rmin" type="float" value="0.0" min="0.0" help="Minimum radial distance."/>
-        <param argument="rmax" type="float" value="10.0" min="0.0" help="Maximum radial distance."/>
+    <xml name="plot_limits_energy">
+        <param name="x_limit_min" type="float" label="Minimum plot energy (eV)" optional="true" help="If set, plot will be limited to this value on the x axis."/>
+        <param name="x_limit_max" type="float" label="Maximum plot energy (eV)" optional="true" help="If set, plot will be limited to this value on the x axis."/>
     </xml>
 </macros>
\ No newline at end of file