--- a/pm_uep_opt.xml	Fri Feb 03 15:40:08 2023 +0000
+++ b/pm_uep_opt.xml	Tue Jul 18 13:27:15 2023 +0000
@@ -1,32 +1,18 @@
-<tool id="pm_uep_opt" name="PyMuonSuite AIRSS UEP Optimise" version="@TOOL_VERSION@+galaxy@WRAPPER_VERSION@" python_template_version="3.5" profile="22.05" license="MIT">
+<tool id="pm_uep_opt" name="PyMuonSuite AIRSS UEP Optimise" version="@PYMUONSUITE_VERSION@+galaxy@WRAPPER_VERSION@" python_template_version="3.5" profile="22.05" license="MIT">
     <description>run UEP optimisation</description>
     <macros>
-        <!-- version of underlying tool (PEP 440) -->
-        <token name="@TOOL_VERSION@">0.2.3</token>
         <!-- version of this tool wrapper (integer) -->
         <token name="@WRAPPER_VERSION@">0</token>
-        <!-- citation should be updated with every underlying tool version -->
-        <!-- typical fields to update are version, month, year, and doi -->
-        <token name="@TOOL_CITATION@">
-            @software{Sturniolo_pymuon-suite_2022,
-                author = {Sturniolo, Simone and Liborio, Leandro and Chadwick, Eli and Murgatroyd, Laura and Laverack, Adam and Mudaraddi, Anish and {Muon Spectroscopy Computational Project}},
-                license = {GPL-3.0},
-                month = {8},
-                title = {{pymuon-suite}},
-                url = {https://github.com/muon-spectroscopy-computational-project/pymuon-suite},
-                version = {v0.2.3},
-                doi = {10.5281/zenodo.7025644},
-                year = {2022}
-            }
-        </token>
+        <import>muon_macros.xml</import>
     </macros>
     <creator>
         <person givenName="Jyothish" familyName="Thomas" identifier="https://orcid.org/0000-0003-4724-6924"/>
         <person givenName="Eli" familyName="Chadwick" url="https://github.com/elichad" identifier="https://orcid.org/0000-0002-0035-6475"/>
+        <person givenName="Patrick" familyName="Austin" url="https://github.com/patrick-austin" identifier="https://orcid.org/0000-0002-6279-7823"/>
         <organization url="https://muon-spectroscopy-computational-project.github.io/index.html" name="The Muon Spectroscopy Computational Project"/>
     </creator>
     <requirements>
-        <requirement type="package" version="@TOOL_VERSION@">pymuonsuite</requirement>
+        <requirement type="package" version="@PYMUONSUITE_VERSION@">pymuonsuite</requirement>
         <requirement type="package" version="3.0">zip</requirement>
     </requirements>
     <required_files>
@@ -34,101 +20,230 @@
         <include type="literal" path="get_out_folder.py"/>
     </required_files>
     <command detect_errors="exit_code"><![CDATA[
-        unzip '$muonated_structures' &&
-        if test -f "params.yaml"; then echo "params.yaml present"; else echo "params.yaml missing" && exit 64; fi &&
-        if ( test -f input_structure.* ) ; then echo "input structure present"; else echo "input structure missing" && exit 64; fi &&
-        ## PyMuonSuite will attempt to load '<chden_seed>.den_fmt' and
-        ## '<chden_seed>.castep', but $charge_density.name may or may not end
-        ## with '.den_fmt' so strip any extension and manually include 'den_fmt'
-        ## and '.castep' in the linked file names
-        charge_density_name=\$(sed 's/ //g' <<< '$charge_density.name') &&
-        charge_density_seed=\${charge_density_name%%.*} &&
-        echo "Calculated seed name \$charge_density_seed" &&
-        ln -s '$charge_density' \$charge_density_seed.den_fmt && 
-        ln -s '$castep_log' \$charge_density_seed.castep &&
+        cat $params_yaml > params.yaml &&
+        structure_name_internal='input_structure.$structure.ext' &&
+        ln -s '$structure' \$structure_name_internal &&
+        pm-muairss -t w \$structure_name_internal params.yaml &&
+        #if $generation_params.allpos_ext != "none":
+            mv allpos.$generation_params.allpos_ext $allpos_file &&
+        #end if
+        ln -s '$charge_density' input_structure.den_fmt && 
+        ln -s '$castep_log' input_structure.castep &&
         out_folder="`python '${__tool_directory__}/get_out_folder.py'`" &&
         find \$out_folder -type f -name "*.yaml" | xargs sed -i "s#^chden_path: .*#chden_path: .\/#g" &&
-        find \$out_folder -type f -name "*.yaml" | xargs sed -i "s#^chden_seed: .*#chden_seed: \$charge_density_seed#g" &&
+        find \$out_folder -type f -name "*.yaml" | xargs sed -i "s#^chden_seed: .*#chden_seed: input_structure#g" &&
         bash '${__tool_directory__}/run.sh' \$out_folder &&
-        zip -r out_zip.zip \$out_folder params.yaml input_structure.* &&
+        zip -r out_zip.zip \$out_folder params.yaml \$structure_name_internal &&
         find \$out_folder > tree.txt
     ]]></command>
+    <configfiles>
+        <configfile name="params_yaml">poisson_r: $generation_params.poisson_r
+vdw_scale: $generation_params.vdw_scale
+#if $generation_params.random_seed:
+random_seed: $generation_params.random_seed
+#end if
+#if $generation_params.allpos_ext!="none":
+allpos_filename: allpos.$generation_params.allpos_ext
+#end if
+calculator: uep
+charged: true
+geom_steps: $optimisation_params.geom_steps
+geom_force_tol: $optimisation_params.geom_force_tol
+uep_gw_factor: $optimisation_params.uep_gw_factor</configfile>
+    </configfiles>
     <inputs>
-        <param type="data" name="muonated_structures" label="Muonated structures (.zip)" format="zip" help="A zip folder containing muonated structures, as generated by the 'Generate muonated structures' tool or pm-muairss."/>
+        <param type="data" name="structure" label="Structure file" format="cell, cif, extxyz" help="The structure to add muons to. The original file will be preserved. Accepted file types: CELL, CIF, XYZ (extended only)."/>
         <param type="data" name="charge_density" label="Charge density file (.den_fmt)" format="den_fmt" help="The charge density file created by CASTEP during your initial DFT simulation for the structure."/>
         <param type="data" name="castep_log" label="CASTEP log (.castep)" format="castep" help="The CASTEP log for your initial DFT simulation for the structure."/>
+        <section name="generation_params" expanded="true" title="Generation Parameters">
+            <param type="float" argument="poisson_r" value="0.8" min="0.0" label="Poisson radius" help=" Poisson sphere radius to use for random generation. No two starting muon positions will be closer than this distance. Smaller values make for bigger structure sets."/>
+            <param type="float" argument="vdw_scale" value="0.5" min="0.0" label="Van der Waals scale" help="Van der Waals scaling factor to use when generating muon sites to avoid existing atoms. Smaller values will allow muons to get closer to the other ions."/>
+            <param type="integer" argument="random_seed" label="Random Seed" optional="true" help="Seed used to randomise the positions of the muons in the generated structures."/>
+            <param type="select" name="allpos_ext" value="none" label="Format for file containing all muon positions (optional)" help="If set, generates an extra output file containing all the muon positions using the specified file type.">
+                <option value="none">Do not generate</option>
+                <option value="cell">CELL</option>
+                <option value="cif">CIF</option>
+                <option value="xyz">XYZ</option>
+            </param>
+        </section>
+        <section name="optimisation_params" expanded="true" title="Optimisation Parameters">
+            <param type="integer" argument="geom_steps" value="30" min="0" label="Geometry optimization steps" help="Maximum number of geometry optimisation steps."/>
+            <param type="float" argument="geom_force_tol" value="0.05" min="0.0" label="Geometry optimization tolerance" help="Tolerance on geometry optimisation in units of eV/Å."/>
+            <param type="float" argument="uep_gw_factor" optional="true" value="5.0" min="0.0" label="Gaussian width factor" help="Gaussian width factor for UEP calculation. Higher values will make the potential of atomic nuclei closer to the point-like approximation but may introduce artifacts."/>
+        </section>
         <param type="hidden" name="testing" label="Test mode" value="false"/>
     </inputs>
     <outputs>
-        <data label="UEP results for $muonated_structures.name" name="uep_results" format="zip" from_work_dir="out_zip.zip"/>
+        <data name="uep_results" format="zip" from_work_dir="out_zip.zip"/>
+        <data label="File containing the positions of all muon sites" name="allpos_file" auto_format="true">
+            <filter>(generation_params["allpos_ext"] != "none")</filter>
+        </data>
         <data label="File tree (testing only)" name="file_tree" format="txt" from_work_dir="tree.txt" hidden="true">
             <filter>(testing == "true")</filter>
         </data>
     </outputs>
     <tests>
-        <test expect_num_outputs="2">
+        <!-- UEP with CELL -->
+        <test expect_num_outputs="3">
             <param name="testing" value="true"/>
-            <param name="muonated_structures" value="muonated.zip" ftype="zip"/>
+            <param name="structure" value="Si.cell" ftype="cell"/>
             <param name="charge_density" value="Si.den_fmt" ftype="den_fmt"/>
             <param name="castep_log" value="Si.castep" ftype="castep"/>
+            <section name="generation_params">
+                <param name="poisson_r" value="1.18"/>
+                <param name="vdw_scale" value="0.25"/>
+                <param name="random_seed" value="7357"/>
+                <param name="allpos_ext" value="cell"/>
+            </section>
+            <section name="optimisation_params">
+                <param name="geom_steps" value="300"/>
+                <param name="uep_gw_factor" value="4.0"/>
+            </section>
+            <output name="allpos_file" file="allpos.cell"/>
             <output name="file_tree" file="tree.txt" ftype="txt" sort="true">
                 <assert_contents>
-                    <has_size value="1780"/>
+                    <has_size value="3280"/>
                 </assert_contents>
             </output>
         </test>
-        <test expect_num_outputs="2">
+        <!-- UEP with XYZ -->
+        <test expect_failure="true">
             <param name="testing" value="true"/>
-            <param name="muonated_structures" value="muonated-default-folder.zip" ftype="zip"/>
+            <param name="structure" value="Si.xyz" ftype="xyz"/>
+            <param name="charge_density" value="Si.den_fmt" ftype="den_fmt"/>
+            <param name="castep_log" value="Si.castep" ftype="castep"/>
+            <section name="generation_params">
+                <param name="poisson_r" value="1.18"/>
+                <param name="vdw_scale" value="0.25"/>
+                <param name="random_seed" value="7357"/>
+            </section>
+            <section name="optimisation_params">
+                <param name="geom_steps" value="300"/>
+                <param name="uep_gw_factor" value="4.0"/>
+            </section>
+            <assert_stderr>
+                <has_text text="TypeError: cannot unpack non-iterable NoneType object"/>
+            </assert_stderr>
+        </test>
+        <!-- UEP with extended XYZ -->
+        <test expect_num_outputs="3">
+            <param name="testing" value="true"/>
+            <param name="structure" value="Si.extxyz" ftype="extxyz"/>
             <param name="charge_density" value="Si.den_fmt" ftype="den_fmt"/>
             <param name="castep_log" value="Si.castep" ftype="castep"/>
-            <output name="file_tree" file="tree-default-folder.txt" sort="true"/>
+            <section name="generation_params">
+                <param name="poisson_r" value="1.18"/>
+                <param name="vdw_scale" value="0.25"/>
+                <param name="random_seed" value="7357"/>
+                <param name="allpos_ext" value="xyz"/>
+            </section>
+            <section name="optimisation_params">
+                <param name="geom_steps" value="300"/>
+                <param name="uep_gw_factor" value="4.0"/>
+            </section>
+            <output name="allpos_file" file="allpos.xyz"/>
+            <output name="file_tree" file="tree.txt" ftype="txt" sort="true">
+                <assert_contents>
+                    <has_size value="3280"/>
+                </assert_contents>
+            </output>
         </test>
+        <!-- UEP with CIF -->
+        <test expect_num_outputs="3">
+            <param name="testing" value="true"/>
+            <param name="structure" value="Si.cif" ftype="cif"/>
+            <param name="charge_density" value="Si.den_fmt" ftype="den_fmt"/>
+            <param name="castep_log" value="Si.castep" ftype="castep"/>
+            <section name="generation_params">
+                <param name="poisson_r" value="1.18"/>
+                <param name="vdw_scale" value="0.25"/>
+                <param name="random_seed" value="7357"/>
+                <param name="allpos_ext" value="cif"/>
+            </section>
+            <section name="optimisation_params">
+                <param name="geom_steps" value="300"/>
+                <param name="uep_gw_factor" value="4.0"/>
+            </section>
+            <output name="allpos_file" file="allpos.cif"/>
+            <output name="file_tree" file="tree.txt" ftype="txt" sort="true">
+                <assert_contents>
+                    <has_size value="3280"/>
+                </assert_contents>
+            </output>
+        </test>
+        <!-- Expect total failure for badly formatted den_fmt -->
         <test expect_failure="true">
             <param name="testing" value="true"/>
-            <param name="muonated_structures" value="muonated-default-folder.zip" ftype="zip"/>
+            <param name="structure" value="Si.cell" ftype="cell"/>
             <param name="castep_log" value="Si.castep" ftype="castep"/>
             <param name="charge_density" value="Sifail.den_fmt" ftype="den_fmt"/>
             <assert_stderr>
                 <has_text text="Error:  FMT file has no header"/>
+                <has_text text="ERROR: All optimisations failed"/>
             </assert_stderr>
         </test>
         <!-- check that file_tree filter works -->
         <test expect_num_outputs="1">
             <param name="testing" value="false"/>
-            <param name="muonated_structures" value="muonated.zip" ftype="zip"/>
+            <param name="structure" value="Si.cell" ftype="cell"/>
             <param name="charge_density" value="Si.den_fmt" ftype="den_fmt"/>
             <param name="castep_log" value="Si.castep" ftype="castep"/>
+            <section name="generation_params">
+                <param name="poisson_r" value="1.18"/>
+                <param name="vdw_scale" value="0.25"/>
+                <param name="random_seed" value="7357"/>
+            </section>
+            <section name="optimisation_params">
+                <param name="geom_steps" value="300"/>
+                <param name="uep_gw_factor" value="4.0"/>
+            </section>
         </test>
         <!-- Check that the chden parameters are set correctly when $charge_density.name does not have the .den_fmt extension -->
         <test expect_num_outputs="2">
             <param name="testing" value="true"/>
-            <param name="muonated_structures" value="muonated.zip" ftype="zip"/>
+            <param name="structure" value="Si.cell" ftype="cell"/>
             <param name="charge_density" value="Si_den_fmt" ftype="den_fmt"/>
             <param name="castep_log" value="Si.castep" ftype="castep"/>
+            <section name="generation_params">
+                <param name="poisson_r" value="1.18"/>
+                <param name="vdw_scale" value="0.25"/>
+                <param name="random_seed" value="7357"/>
+            </section>
+            <section name="optimisation_params">
+                <param name="geom_steps" value="300"/>
+                <param name="uep_gw_factor" value="4.0"/>
+            </section>
             <output name="file_tree" file="tree.txt" ftype="txt" sort="true">
                 <assert_contents>
-                    <has_size value="1780"/>
+                    <has_size value="3280"/>
                 </assert_contents>
             </output>
         </test>
     </tests>
     <help><![CDATA[
-        Given a set of input muonated structures (structures containing a muon), applies UEP optimisation to refine them.
+        Given an input structure, this tool generates a set of duplicate structures,
+        each with a single muon added in a random location.
+        These are known as "muonated" structures.
+        Following this, UEP optimisation is used to refine the positions of the muons.
 
-        This tool requires CASTEP inputs for the UEP method, meaning your initial DFT simulation must have been done with CASTEP.
-        The structure you used as input to the '`Generate muonated structures`_' tool or pm-muairss beforehand should come from the same simulation run.
+        This tool requires CASTEP inputs in order to perform the UEP method,
+        meaning your initial DFT simulation must have been done with CASTEP.
+        This simulation should have used the same structure as is used here.
 
-        Command-line usage: pm-uep-opt [-h] [-t {r,w}] structures parameter_file
+        Following this, the individual refined positions can be clustered to
+        represent potential stopping sites by using '`PyMuonSuite AIRSS Cluster`_'.
 
-        .. _Generate muonated structures: /tool_runner?tool_id=pm_muairss_write
+        Command-line usage (generation): pm-muairss [-h] -t w structures parameter_file
+        
+        Command-line usage (optimisation): pm-uep-opt [-h] parameter_file
 
-        PyMuonSuite is distributed under the GPLv3 license. This tool wrapper is distributed under the MIT license.
+        .. _PyMuonSuite AIRSS Cluster: /tool_runner?tool_id=pm_muairss_read
+
+        PyMuonSuite is distributed under the GPLv3 license.
+        This tool wrapper is distributed under the MIT license.
     ]]></help>
     <citations>
-        <citation type="bibtex">
-            @TOOL_CITATION@
-        </citation>
+        <citation type="doi">@PYMUONSUITE_CITATION@</citation>
         <citation type="doi">10.1063/1.5024450</citation>
         <citation type="doi">10.1063/1.5085197</citation>
         <citation type="doi">10.1063/5.0012381</citation>