Merge pull request #20 from N3PDF/add_autolink

Add an autolink action

doc/source/installation.rst

@@ -28,10 +28,13 @@ if you are planning to extend or develop code use instead:
   pip install -e .
 
 
-MG5_aMC\@NLO
------------
+.. _plugin-label:
+
+MG5_aMC\@NLO Plugin
+--------------------
 
 A valid installation of MG5_aMC\@NLO (2.8+) is necessary in order to generate matrix elements.
+
 If you already have a valid installation, please add the following environment variable pointing to the right directory: ``MADGRAPH_PATH``.
 Below are the instructions for MG5_aMC\@NLO 3.1.0, for a more recent release please visit the MG5_aMC\@NLO `site <https://launchpad.net/mg5amcnlo>`_.
 
@@ -42,12 +45,16 @@ Below are the instructions for MG5_aMC\@NLO 3.1.0, for a more recent release ple
   export MADGRAPH_PATH=${PWD}/MG5_aMC_v3_1_0
 
 
-.. _plugin-label:
+Once MG5_aMC\@NLO is installed, all that's left is to link the ``madflow`` plugin inside
+the MG5_aMC\@NLO folder.
+
+
+.. code-block:: bash
+
+  madflow --autolink
 
-Plugin
-------
 
-In order to install the ``madflow`` plugin in MG5_aMC\@NLO, it is necessary to link the
+If you prefer to link the plugin manually, it is necessary to link the
 ``madgraph_plugin`` folder inside the ``PLUGIN`` directory of MG5_aMC\@NLO.
 For instance, if the environment variable ``$MADGRAPH_PATH`` is pointing to the MG5_aMC root:
 

python_package/madflow/config.py

@@ -82,22 +82,25 @@ def complex_me(cmp):
 def get_madgraph_path(madpath=None):
     """Return the path to the madgrapt root"""
     if madpath is None:
-        madpath = os.environ.get("MADGRAPH_PATH", "../../../mg5amcnlo")
+        madpath = os.environ.get("MADGRAPH_PATH", "mg5amcnlo")
     madgraph_path = Path(madpath)
     if not madgraph_path.exists():
         raise ValueError(
             f"{madgraph_path} does not exist. "
             "Needs a valid path for Madgraph, can be given as env. variable MADGRAPH_PATH"
         )
+    # If the path exists, check whether the madgraph executable is there
+    _ = get_madgraph_exe(madgraph_path)
     return madgraph_path
 
 
 def get_madgraph_exe(madpath=None):
     """Return the path to the madgraph executable"""
-    madpath = get_madgraph_path(madpath)
+    if madpath is None:
+        madpath = get_madgraph_path(madpath)
     mg5_exe = madpath / "bin/mg5_aMC"
     if not mg5_exe.exists():
-        raise ValueError(f"Madgraph executablec ould not be found at {mg5_exe}")
+        raise ValueError(f"Madgraph executable could not be found at {mg5_exe}")
     return mg5_exe
 
 

python_package/madflow/scripts/madflow_exec.py

@@ -20,6 +20,9 @@
 import sys
 import time
 import shutil
+import tarfile
+import requests
+import datetime
 import itertools
 import importlib
 import argparse
@@ -40,7 +43,6 @@
     guess_events_limit,
 )
 from madflow.phasespace import PhaseSpaceGenerator
-from madflow.lhe_writer import LheWriter
 
 from vegasflow import VegasFlow
 from pdfflow import mkPDF
@@ -151,8 +153,73 @@ def _generate_initial_states(matrices):
     return initial_flavours
 
 
+def _autolink(madpath):
+    """Links the madflow madgraph plugin into the MG5_aMC directory
+    If needed it downloads the plugin from the right github directory
+    """
+    try:
+        madgraph_path = get_madgraph_path(madpath)
+    except ValueError:
+        print(
+            """> The madgraph path could not be autodiscovered
+> Please, set the MADGRAPH_PATH to the madgraph root directory
+> Or add the exact path to --autolink (e.g. madflow --autolink /home/mg5)
+and don't forget to set `export MADGRAPH_PATH=/path/to/madgraph` in your favourite .rc file!"""
+        )
+        sys.exit(0)
+
+    # Check whether we already have a plugin there
+    plugin_path = madgraph_path / "PLUGIN/pyout"
+    if plugin_path.exists():
+        print("The plugin folder already exists")
+        yn = input("Do you want to remove it and link the new one? ")
+        if yn.lower() != "y" and yn.lower() != "yes":
+            sys.exit(0)
+        # Don't fully remove it, just move it around
+        new_path = plugin_path
+        nn = 0
+        while new_path.exists():
+            today_name = datetime.datetime.now().strftime(f".backup-%d-%m-%y-n{nn}")
+            new_path = new_path.with_suffix(today_name)
+            nn += 1
+        plugin_path.rename(new_path)
+
+    # If this is a develop setup, link the repository version
+    test_path = Path(__file__).parent / "../../../madgraph_plugin"
+    if test_path.exists():
+        plugin_path.symlink_to(test_path)
+    else:
+        # Download plugin
+        latest_plugin = (
+            "https://github.com/N3PDF/madflow/releases/latest/download/madgraph_plugin.tar.gz"
+        )
+        target_path = Path("/tmp/madgraph_plugin.tar.gz")
+        print(f"Downloading plugin from github repository and untaring to {plugin_path}")
+        response = requests.get(latest_plugin, stream=True)
+        if response.status_code == 200:
+            target_path.write_bytes(response.raw.read())
+        with tarfile.open(target_path) as tar:
+            tar.extractall(plugin_path.parent)
+
+    print("Linking finished, exiting")
+
+
+class _MadFlowAutolink(argparse.Action):
+    """Wrapper action around _autolink"""
+
+    def __init__(self, **kw):
+        super().__init__(nargs="?", **kw)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        _autolink(values)
+        parser.exit(0)
+
+
 def madflow_main(args=None, quick_return=False):
     arger = argparse.ArgumentParser(__doc__)
+
+    arger.add_argument("--autolink", help="Link madflow with madgraph", action=_MadFlowAutolink)
+
     arger.add_argument("-v", "--verbose", help="Print extra info", action="store_true")
     arger.add_argument("-p", "--pdf", help="PDF set", type=str, default=DEFAULT_PDF)
     arger.add_argument(
@@ -203,12 +270,21 @@ def madflow_main(args=None, quick_return=False):
     arger.add_argument(
         "--dry_run", help="Generate the madgraph output but don't run anything", action="store_true"
     )
-    arger.add_argument("--events_per_iteration", help="How many events to run per iteration", type=int, default=int(1e6))
+    arger.add_argument(
+        "--events_per_iteration",
+        help="How many events to run per iteration",
+        type=int,
+        default=int(1e6),
+    )
 
     args = arger.parse_args(args)
+
     if quick_return:
         return args, None, None
 
+    # LheWriter needs to be imported after --autolink
+    from madflow.lhe_writer import LheWriter
+
     if args.output is None:
         output_path = Path(tempfile.mkdtemp(prefix="mad_"))
     else:
@@ -398,7 +474,7 @@ def cross_section(xrand, n_dim=ndim, weight=1.0):
             res, err = vegas.run_integration(final_iterations)
             lhe_writer.store_result((res, err))
             proc_folder = output_path / f"Events/{proc_name}"
-            filename = proc_folder / 'cross_err.txt'
+            filename = proc_folder / "cross_err.txt"
             lhe_writer.dump_result(filename)
             logger.info("Written LHE file to %s", proc_folder)
     else:

readme.md

@@ -51,12 +51,14 @@ You can check your installed version of `pdfflow` with: `python -c 'import pdffl
 ## Install plugin in MG5_aMC
 
 In order to install the `madflow` plugin in MG5_aMC@NLO, it is necessary to link the `madgraph_plugin` folder inside the `PLUGIN` directory of MG5_aMC@NLO.
-For instance, if the environment variable `$MADGRAPH_PATH` is pointing to the MG5_aMC root:
+For instance, if the environment variable `$MADGRAPH_PATH` is pointing to the MG5_aMC root and you are currently in the repository root.
 
 ```bash
     ln -s ${PWD}/madgraph_plugin ${MADGRAPH_PATH}/PLUGIN/pyout
 ```
 
+The link can be performed automagically with the `madflow --autolink` option.
+
 ## Use `madflow`
 
 For a more precise description of what `madflow` can do please visit the online documentation.

setup.py

@@ -8,7 +8,7 @@
 import os
 from setuptools import setup, find_packages
 
-requirements = ["vegasflow", "pdfflow"]
+requirements = ["vegasflow", "pdfflow", "requests"]
 # soft-requirements due to vegasflow and pdfflow are:
 #   tensorflow, joblib, numpy
 package_name = "madflow"