ReflGenerateNotebook.cpp

#include "MantidQtCustomInterfaces/ReflGenerateNotebook.h"
#include "MantidAPI/NotebookWriter.h"
#include "MantidQtCustomInterfaces/ParseKeyValueString.h"

#include <sstream>
#include <fstream>
#include <memory>
#include <boost/tokenizer.hpp>
#include <boost/regex.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/foreach.hpp>
#include <boost/range/combine.hpp>

namespace MantidQt {
  namespace CustomInterfaces {

    ReflGenerateNotebook::ReflGenerateNotebook(std::string name, QReflTableModel_sptr model,
                                               const std::string instrument, const int runs_column,
                                               const int transmission_column, const int options_column,
                                               const int angle_column, const int min_q, const int max_q,
                                               const int d_qq, const int scale_column, const int group_column) :
      m_wsName(name), m_model(model), m_instrument(instrument),
      col_nums{runs_column, transmission_column, options_column, angle_column, min_q, max_q, d_qq, scale_column, group_column} { }

    /**
      Generate an ipython notebook
      @param groups : groups of rows which were stitched
      @returns ipython notebook string
      */
    std::string ReflGenerateNotebook::generateNotebook(std::map<int, std::set<int>> groups, std::set<int> rows) {
      std::unique_ptr<Mantid::API::NotebookWriter> notebook(new Mantid::API::NotebookWriter());

      const std::string plotFunctionsTitle = "Plot functions\n---------------";
      notebook->markdownCell(plotFunctionsTitle);
      notebook->codeCell(plotsFunctionString());

      std::string title_string;
      if (!m_wsName.empty()) {
        title_string = "Processed data from workspace: " + m_wsName + "\n---------------------";
      }
      else {
        title_string = "Processed data\n---------------------";
      }
      title_string += "\nNotebook generated from the ISIS Reflectometry (Polref) Interface";
      notebook->markdownCell(title_string);

      notebook->markdownCell(tableString(m_model, col_nums, rows));

      int groupNo = 1;
      for (auto gIt = groups.begin(); gIt != groups.end(); ++gIt, ++groupNo) {
        const std::set<int> groupRows = gIt->second;

        // Announce the stitch group in the notebook
        notebook->markdownCell(
          "Stitch group " + static_cast<std::ostringstream*>( &(std::ostringstream() << groupNo) )->str());

        //Reduce each row
        std::ostringstream code_string;
        std::tuple<std::string, std::string, std::string> reduce_row_string;
        std::vector<std::string> unstitched_ws;
        std::vector<std::string> IvsLam_ws;
        code_string << "#Load and reduce\n";
        for (auto rIt = groupRows.begin(); rIt != groupRows.end(); ++rIt) {
          reduce_row_string = reduceRowString(*rIt, m_instrument, m_model, col_nums);
          code_string << std::get<0>(reduce_row_string);
          unstitched_ws.push_back(std::get<1>(reduce_row_string));
          IvsLam_ws.push_back(std::get<2>(reduce_row_string));
        }
        notebook->codeCell(code_string.str());

        // Stitch group
        std::tuple<std::string, std::string> stitch_string = stitchGroupString(groupRows, m_instrument, m_model, col_nums);
        notebook->codeCell(std::get<0>(stitch_string));

        // Group workspaces which should be plotted on same axes
        std::ostringstream plot_string;
        plot_string << "#Group workspaces to be plotted on same axes\n";
        plot_string << "unstitchedGroupWS = GroupWorkspaces(" << vectorParamString("InputWorkspaces", unstitched_ws) << ")\n";
        plot_string << "IvsLamGroupWS = GroupWorkspaces(" << vectorParamString("InputWorkspaces", IvsLam_ws) << ")\n";

        // Plot I vs Q and I vs Lambda graphs
        plot_string << "#Plot workspaces\n";
        std::vector<std::string> workspaceList;
        workspaceList.push_back("unstitchedGroupWS");
        workspaceList.push_back(std::get<1>(stitch_string));
        workspaceList.push_back("IvsLamGroupWS");

        plot_string << plot1DString(workspaceList, "['I vs Q Unstitched', 'I vs Q Stitiched', 'I vs Lambda']");
        notebook->codeCell(plot_string.str());
      }

      return notebook->writeNotebook();
    }

    std::string tableString(QReflTableModel_sptr model, ColNumbers col_nums, const std::set<int> & rows)
    {
      std::ostringstream table_string;

      table_string << "Run(s) | Angle | Transmission Run(s) | Q min | Q max | dQ/Q | Scale | Group | Options\n";
      table_string << "------ | ----- | ------------------- | ----- | ----- | ---- | ----- | ----- | -------\n";

      for(auto rowIt = rows.begin(); rowIt != rows.end(); ++rowIt) {
        table_string << model->data(model->index(*rowIt, col_nums.runs)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.angle)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.transmission)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.qmin)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.qmax)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.dqq)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.scale)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.group)).toString().toStdString() << " | ";
        table_string << model->data(model->index(*rowIt, col_nums.options)).toString().toStdString() << "\n";
      }

      return table_string.str();
    }

    std::string plotsFunctionString()
    {
      return "def plotWithOptions(ax, ws, ops, n):\n"
        "    \"\"\"\n"
        "    Enable/disable legend, grid, limits according to\n"
        "    options (ops) for the given axes (ax).\n"
        "    Plot with or without errorbars.\n"
        "    \"\"\"\n"
        "    ws_plot = ConvertToPointData(ws)\n"
        "    if ops['errorbars']:\n"
        "        ax.errorbar(ws_plot.readX(0), ws_plot.readY(0), yerr=ws_plot.readE(0), label=ws.name())\n"
        "    else:\n"
        "        ax.plot(ws_plot.readX(0), ws_plot.readY(0), label=ws.name())\n"
        "    \n"
        "    ax.grid(ops['grid'])\n"
        "    ax.set_xscale(ops['xScale']); ax.set_yscale(ops['yScale'])\n"
        "    if ops['xLimits'] != 'auto': ax.set_xlim(ops['xLimits'])\n"
        "    if ops['yLimits'] != 'auto': ax.set_ylim(ops['yLimits'])\n"
        "    \n"
        "    # If a list of titles was given, use it to title each subplot\n"
        "    if hasattr(ops['title'], \"__iter__\"):\n"
        "        ax.set_title(ops['title'][n])\n"
        "    if ops['legend'] and hasattr(ops['legendLocation'], \"__iter__\"):\n"
        "        ax.legend(loc=ops['legendLocation'][n])\n"
        "    elif ops['legend']:\n"
        "        ax.legend(loc=ops['legendLocation'])"
        "    \n"
        "\n"
        "def plots(listOfWorkspaces, *args, **kwargs):\n"
        "    \"\"\"\n"
        "    Draw a default reflectivity plot.\n"
        "    Workspaces within a group workspace are plotted together on the same axes.\n"
        "\n"
        "    Examples:\n"
        "    plots(rr)\n"
        "    plots(rr, 'TheGraphTitle')\n"
        "    plots(rr, 'TheGraphTitle', grid=True, legend=True, xScale='linear', yScale='log', xLimits=[0.008, 0.16])\n"
        "    plots(rr, sharedAxes = False, xLimits = [0, 0.1], yLimits = [1e-5, 2], Title='ASF070_07 I=1A T=3K dq/q=2%', legend=True, legendLocation=3, errorbars=False)\n"
        "    \"\"\"\n"
        "\n"
        "    if not hasattr(listOfWorkspaces, \"__iter__\"):\n"
        "        listOfWorkspaces = [listOfWorkspaces]\n"
        "\n"
        "    # Process the function arguments.  In either case(named or unnamed) build a dictionary of options)\n"
        "    keylist = ['title', 'grid', 'legend', 'legendLocation', 'xScale', 'yScale', 'xLimits', 'yLimits', 'sharedAxes', 'errorbars']\n"
        "    defaultValues = ['', True, True, 1, 'log', 'log', 'auto', 'auto', True, 'True']\n"
        "\n"
        "    # Fill ops with the default values\n"
        "    ops=dict(zip(keylist,defaultValues))\n"
        "    for i in range(len(args)):  # copy in values provided in args\n"
        "        defaultValues[i]=args[i]\n"
        "    ops=dict(zip(keylist,defaultValues))\n"
        "\n"
        "    for k in ops.keys():  # copy in any key word given arguments\n"
        "        ops[k]= kwargs.get(k,ops[k])\n"
        "\n"
        "    # Create subplots for workspaces in the list\n"
        "    fig, ax = plt.subplots(1, len(listOfWorkspaces), sharey=ops['sharedAxes'], figsize=(6*len(listOfWorkspaces),4))\n"
        "\n"
        "    if not hasattr(ax, \"__iter__\"):\n"
        "        ax = [ax]\n"
        "\n"
        "    for n, ws in enumerate(listOfWorkspaces):\n"
        "        if type(ws) == mantid.api._api.WorkspaceGroup:\n"
        "            # Plot grouped workspaces on the same axes\n"
        "            for sub_ws in ws:\n"
        "                plotWithOptions(ax[n], sub_ws, ops, n)\n"
        "        else:\n"
        "            plotWithOptions(ax[n], ws, ops, n)\n"
        "            \n"
        "    # If a single title was given, use it to title the whole figure\n"
        "    if not hasattr(ops['title'], \"__iter__\"):\n"
        "        fig.suptitle(ops['title'])\n"
        "    plt.show()\n"
        "    \n"
        "    return plt.gcf()";
    }

    /**
      Create string of python code to stitch workspaces in the same group
      @param rows : rows in the stitch group
      @return tuple containing the python code string and the output workspace name
      */
    std::tuple<std::string, std::string> stitchGroupString(const std::set<int> & rows, const std::string & instrument,
                                                           QReflTableModel_sptr model, ColNumbers col_nums)
    {
      std::ostringstream stitch_string;

      stitch_string << "#Stitch workspaces\n";

      //If we can get away with doing nothing, do.
      if(rows.size() < 2)
        return std::make_tuple("", "");

      //Properties for Stitch1DMany
      std::vector<std::string> workspaceNames;
      std::vector<std::string> runs;

      std::vector<double> params;
      std::vector<double> startOverlaps;
      std::vector<double> endOverlaps;

      //Go through each row and prepare the properties
      for(auto rowIt = rows.begin(); rowIt != rows.end(); ++rowIt)
      {
        const std::string  runStr = model->data(model->index(*rowIt, col_nums.runs)).toString().toStdString();
        const double         qmin = model->data(model->index(*rowIt, col_nums.qmin)).toDouble();
        const double         qmax = model->data(model->index(*rowIt, col_nums.qmax)).toDouble();

        const std::tuple<std::string, std::string> load_ws_string = loadWorkspaceString(runStr, instrument);

        const std::string runNo = getRunNumber(std::get<1>(load_ws_string));
        runs.push_back(runNo);
        workspaceNames.push_back("IvsQ_" + runNo);

        startOverlaps.push_back(qmin);
        endOverlaps.push_back(qmax);
      }

      double dqq = model->data(model->index(*(rows.begin()), col_nums.dqq)).toDouble();

      //params are qmin, -dqq, qmax for the final output
      params.push_back(*std::min_element(startOverlaps.begin(), startOverlaps.end()));
      params.push_back(-dqq);
      params.push_back(*std::max_element(endOverlaps.begin(), endOverlaps.end()));

      //startOverlaps and endOverlaps need to be slightly offset from each other
      //See usage examples of Stitch1DMany to see why we discard first qmin and last qmax
      startOverlaps.erase(startOverlaps.begin());
      endOverlaps.pop_back();

      std::string outputWSName = "IvsQ_" + boost::algorithm::join(runs, "_");

      stitch_string << outputWSName << ", _ = Stitch1DMany(";
      stitch_string << vectorParamString("InputWorkspaces", workspaceNames);
      stitch_string << ", ";
      stitch_string << vectorParamString("Params", params);
      stitch_string << ", ";
      stitch_string << vectorParamString("StartOverlaps", startOverlaps);
      stitch_string << ", ";
      stitch_string << vectorParamString("EndOverlaps", endOverlaps);
      stitch_string << ")\n";

      return std::make_tuple(stitch_string.str(), outputWSName);
    }

    /**
      Create string of comma separated list of parameter values from a vector
      @param param_name : name of the parameter we are creating a list of
      @param param_vec : vector of parameter values
      @return string of comma separated list of parameter values
      */
    template<typename T, typename A>
    std::string vectorParamString(const std::string & param_name, const std::vector<T,A> &param_vec)
    {
      std::ostringstream param_vector_string;

      param_vector_string << param_name << " = '";
      param_vector_string << vectorString(param_vec);
      param_vector_string << "'";

      return param_vector_string.str();
    }

    template<typename T, typename A>
    std::string vectorString(const std::vector<T,A> &param_vec)
    {
      std::ostringstream vector_string;
      const char* separator = "";
      for(auto paramIt = param_vec.begin(); paramIt != param_vec.end(); ++paramIt)
      {
        vector_string << separator << *paramIt;
        separator = ", ";
      }

      return vector_string.str();
    }

    /**
      Create string of python code to create 1D plots from workspaces
      @param ws_names : vector of workspace names to plot
      @return string  of python code to plot I vs Q
      */
    std::string plot1DString(const std::vector<std::string> & ws_names,
                             const std::string & title) {

      std::ostringstream plot_string;

      plot_string << "fig = plots([" << vectorString(ws_names) << "], title=" << title
                  << ", legendLocation=[1, 1, 4])\n";

      return plot_string.str();
    }

    /**
     Create string of python code to run reduction algorithm on the specified row
     @param rowNo : the row in the model to run the reduction algorithm on
     @return tuple containing the python string and the output workspace name
    */
    std::tuple<std::string, std::string, std::string>
    reduceRowString(const int rowNo, const std::string & instrument, QReflTableModel_sptr model, ColNumbers col_nums) {
      std::ostringstream code_string;

      const std::string runStr = model->data(model->index(rowNo, col_nums.runs)).toString().toStdString();
      const std::string transStr = model->data(model->index(rowNo, col_nums.transmission)).toString().toStdString();
      const std::string options = model->data(model->index(rowNo, col_nums.options)).toString().toStdString();

      double theta = 0;

      const bool thetaGiven = !model->data(model->index(rowNo, col_nums.angle)).toString().isEmpty();

      if (thetaGiven)
        theta = model->data(model->index(rowNo, col_nums.angle)).toDouble();

      const std::tuple<std::string, std::string> load_ws_string = loadWorkspaceString(runStr, instrument);
      code_string << std::get<0>(load_ws_string);

      const std::string runNo = getRunNumber(std::get<1>(load_ws_string));
      const std::string IvsLamName = "IvsLam_" + runNo;
      const std::string thetaName = "theta_" + runNo;

      if (!transStr.empty()) {
        const std::tuple<std::string, std::string> trans_string = transWSString(transStr, instrument);
        code_string << std::get<0>(trans_string);
        code_string << "IvsQ_" << runNo << ", " << IvsLamName << ", " << thetaName << " = ";
        code_string << "ReflectometryReductionOneAuto(InputWorkspace = '" << std::get<1>(load_ws_string) << "'";
        code_string << ", " << "FirstTransmissionRun = '" << std::get<1>(trans_string) << "'";
      }
      else {
        code_string << "IvsQ_" << runNo << ", " << IvsLamName << ", " << thetaName << " = ";
        code_string << "ReflectometryReductionOneAuto(InputWorkspace = '" << std::get<1>(load_ws_string) << "'";
      }

      std::string thetaStr;
      if (thetaGiven) {
        code_string << ", " << "ThetaIn = " << theta;
        thetaStr = static_cast<std::ostringstream *>( &(std::ostringstream() << theta))->str();
      }
      else {
        thetaStr = "theta_" + runNo; // Use variable name if we don't have the value
      }

      //Parse and set any user-specified options
      auto optionsMap = parseKeyValueString(options);
      for (auto kvp = optionsMap.begin(); kvp != optionsMap.end(); ++kvp) {
        code_string << ", " << kvp->first << " = " << kvp->second;
      }
      code_string << ")\n";

      const double scale = model->data(model->index(rowNo, col_nums.scale)).toDouble();
      if(scale != 1.0) {
        const std::tuple<std::string, std::string> scale_string = scaleString(runNo, scale);
        code_string << std::get<0>(scale_string);
      }

      const std::tuple<std::string, std::string> rebin_string = rebinString(rowNo, runNo, model, col_nums);
      code_string << std::get<0>(rebin_string);

      return std::make_tuple(code_string.str(), std::get<1>(rebin_string), IvsLamName);
    }

     /**
      Create string of python code to run the scale algorithm
      @param runNo : the number of the run to scale
      @param scale : value of scaling factor to use
      @return tuple of strings of python code and output workspace name
      */
    std::tuple<std::string, std::string>
     scaleString(const std::string & runNo, const double scale)
    {
      std::ostringstream scale_string;

      scale_string << "IvsQ_" << runNo << " = Scale(";
      scale_string << "InputWorkspace = IvsQ_" << runNo;
      scale_string << ", Factor = " << 1.0 / scale;
      scale_string << ")\n";

      return std::make_tuple(scale_string.str(), "IvsQ" + runNo);
    }

    /**
      Create string of python code to convert to point data, which can be plotted
      @param wsName : name of workspace to convert to point data
      @return tuple of strings of python code and output workspace name
      */
    std::tuple<std::string, std::string> convertToPointString(const std::string & wsName)
    {
      const std::string output_name = wsName + "_plot";
      std::ostringstream convert_string;
      convert_string << output_name << " = ConvertToPointData(" << wsName << ")\n";

      return std::make_tuple(convert_string.str(), output_name);
    }

    /**
     Create string of python code to rebin data in a workspace
     @param rowNo : the number of the row to rebin
     @param runNo : the number of the run to rebin
     @return tuple of strings of python code and output workspace name
    */
    std::tuple<std::string, std::string>
      rebinString(const int rowNo, const std::string & runNo, QReflTableModel_sptr model, ColNumbers col_nums)
    {
      //We need to make sure that qmin and qmax are respected, so we rebin to
      //those limits here.
      std::ostringstream rebin_string;
      rebin_string << "IvsQ_" << runNo << " = ";
      rebin_string << "Rebin(";
      rebin_string << "IvsQ_" << runNo;

      const double qmin = model->data(model->index(rowNo, col_nums.qmin)).toDouble();
      const double qmax = model->data(model->index(rowNo, col_nums.qmax)).toDouble();
      const double dqq = model->data(model->index(rowNo, col_nums.dqq)).toDouble();

      rebin_string << ", " << "Params = ";
      rebin_string << "'" << qmin << ", " << -dqq << ", " << qmax << "'";
      rebin_string << ")\n";

      return std::make_tuple(rebin_string.str(), "IvsQ_" + runNo);
    }

    /**
     Create string of python code to create a transmission workspace
     @param trans_ws_str : string of workspaces to create transmission workspace from
     @return tuple of strings of python code and output workspace name
    */
    std::tuple<std::string, std::string> transWSString(const std::string & trans_ws_str, const std::string & instrument)
    {
      const size_t maxTransWS = 2;

      std::vector<std::string> transVec;
      std::ostringstream trans_string;
      std::vector<std::string> trans_ws_name;

      //Take the first two run numbers
      boost::split(transVec, trans_ws_str, boost::is_any_of(","));
      if(transVec.size() > maxTransWS)
        transVec.resize(maxTransWS);

      std::tuple<std::string, std::string> load_tuple;
      for(auto it = transVec.begin(); it != transVec.end(); ++it)
        load_tuple = loadWorkspaceString(*it, instrument);
        trans_ws_name.push_back(std::get<1>(load_tuple));
        trans_string << std::get<0>(load_tuple);

      //The runs are loaded, so we can create a TransWS
      std::string wsName = "TRANS_" + getRunNumber(trans_ws_name[0]);
      if(trans_ws_name.size() > 1)
        wsName += "_" + getRunNumber(trans_ws_name[1]);
      trans_string << wsName << " = ";
      trans_string << "CreateTransmissionWorkspaceAuto(";
      trans_string << "FirstTransmissionRun = '" << trans_ws_name[0] << "'";
      if(trans_ws_name.size() > 1)
        trans_string << ", SecondTransmissionRun = '" << trans_ws_name[1] << "'";
      trans_string << ")\n";

      return std::make_tuple(trans_string.str(), wsName);
    }

    /**
     Get run number from workspace name
     @param ws_name : workspace name
     @return run number, as a string
    */
    std::string getRunNumber(const std::string & ws_name) {
      //Matches TOF_13460 -> 13460
      boost::regex outputRegex("(TOF|IvsQ|IvsLam)_([0-9]+)");

      //Matches INTER13460 -> 13460
      boost::regex instrumentRegex("[a-zA-Z]{3,}([0-9]{3,})");

      boost::smatch matches;

      if (boost::regex_match(ws_name, matches, outputRegex)) {
        return matches[2].str();
      }
      else if (boost::regex_match(ws_name, matches, instrumentRegex)) {
        return matches[1].str();
      }

      //Resort to using the workspace name
      return ws_name;
    }

    /**
     Create string of python code to load workspaces
     @param runStr : string of workspaces to load
     @return tuple of strings of python code and output workspace name
    */
    std::tuple<std::string, std::string> loadWorkspaceString(const std::string & runStr, const std::string & instrument) {
      std::vector<std::string> runs;
      boost::split(runs, runStr, boost::is_any_of("+"));

      std::ostringstream load_strings;

      //Remove leading/trailing whitespace from each run
      for (auto runIt = runs.begin(); runIt != runs.end(); ++runIt)
        boost::trim(*runIt);

      const std::string outputName = "TOF_" + boost::algorithm::join(runs, "_");

      std::tuple<std::string, std::string> load_string;

      load_string = loadRunString(runs[0], instrument);
      load_strings << std::get<0>(load_string);

      // EXIT POINT if there is only one run
      if(runs.size() == 1) {
        return std::make_tuple(load_strings.str(), std::get<1>(load_string));
      }
      load_strings << outputName << " = " << std::get<1>(load_string) << "\n";

      // Load each subsequent run and add it to the first run
      for(auto runIt = std::next(runs.begin()); runIt != runs.end(); ++runIt)
      {
        load_string = loadRunString(*runIt, instrument);
        load_strings << std::get<0>(load_string);
        load_strings << plusString(std::get<1>(load_string), outputName);
      }

      return std::make_tuple(load_strings.str(), outputName);
    }

    /**
     Create string of python code to run the Plus algorithm on specified workspaces
     @param input_name : name of workspace to add to the other workspace
     @param output_name : other workspace will be added to the one with this name
     @return string of python code
    */
    std::string plusString(const std::string & input_name, const std::string & output_name)
    {
      std::ostringstream plus_string;

      plus_string << output_name << " = Plus('LHSWorkspace' = " << output_name;
      plus_string << ", 'RHSWorkspace' = " << input_name;
      plus_string << ")\n";

      return plus_string.str();
    }

    /**
     Create string of python code to load a single workspace
     @param run : run to load
     @return tuple of strings of python code and output workspace name
    */
    std::tuple<std::string, std::string> loadRunString(const std::string & run, const std::string & instrument) {
      std::ostringstream load_string;
      // We do not have access to AnalysisDataService from notebook, so must load run from file
      const std::string filename = instrument + run;
      const std::string ws_name = "TOF_" + run;
      load_string << ws_name << " = ";
      load_string << "Load(";
      load_string << "Filename = '" << filename << "'";
      load_string << ")\n";

      return std::make_tuple(load_string.str(), ws_name);
    }

  }
}