Tutorial: How to write an Ot2Rec plugin

Guide to writing plugins for Ot2Rec v0.2.x

Introduction

Plugins in Ot2Rec achieve four main objectives:

1. Capture user arguments for the task. (These are saved in a yaml file)
2. Generate the correct commands to call other programs to complete the task, e.g. IMOD
3. Save the metadata along the way
4. (Optionally) add a section to the report showing how the task performed.

This guide will show you how to write an Ot2Rec plugin called plugin.

To-Do list

• Add the user arguments to magicgui.py
• Add the desired layout of plugin.yaml to params.py
• Create plugin.py, which holds most of the logic behind your new plugin.
• Add your command-line binding to setup.py. This binding allows users to call your plugin straight from the terminal.

Step 0: Setup your conda environment

Clone Ot2Rec from Github:

git clone https://github.com/rosalindfranklininstitute/Ot2Rec.git

Create a new conda environment and install Ot2Rec in editable mode:

conda create -n ot2rec
conda activate ot2rec
pip install -e .

This should update your installation of Ot2Rec when changes are made so you are always running the latest version.

Step 1: Capture user arguments

We use magicgui to generate Ot2Rec's GUI to capture user arguments. These argument capture functions are added to the magicgui.py file. One GUI is created for each plugin, an example is shown below:

add image here

Here is a minimal example of the magicgui functions we need for our plugin.

@mg(
    call_button="Create config file",
    layout="vertical",
    result_widget=False,

    project_name={"label": "Project name *"},
    pixel_size={"label": "Pixel size in A",
                "min": 0.001,
    },
    rootname={"label": "Rootname of current project (required if different from project name"},
    suffix={"label": "Suffix of project files"},
    input_mrc_folder={
        "label": "Folder containing input mrc's",
        "mode": "d",
    },
    output_path={
        "label": "Path to output folder",
        "mode": "d",
    },
    choice_abc={
        "label": "Choice between a, b, c",
        "choices": ["a", "b", "c]
    },
)
def get_args_plugin(
        project_name="",
        pixel_size=0.00,
        rootname="",
        suffix="",
        input_mrc_folder=Path("./plugin_input"),
        output_path=Path("./plugin_output"),
        recon_algo="WBP",
):
    return locals()

First, we use the magicgui decorator @mg to create our widgets. The widgets we have on all our plugins are:

• project_name | name of the dataset we are processing
• rootname | prefix of all the files we want to work with, defaults to the project_name if not specified.
• suffix | suffix added to the end of the filenames, can be left as None.
• some sort of input, e.g., input_mrc_folder
• an output path, e.g., output_path

Step 2: Create template of yaml to hold arguments

The plugin.yaml file holds the information needed to set up the task to be performed, e.g., input filepaths, parameters, output filepaths. Some of these are captured from the magicgui in the previous section, but others can be calculated or populated with default values.

The yaml file is first generated from a template in params.py, its values are then populated by code in plugin.py with user arguments captured from magicgui.py.

An example yaml file for our plugin looks like:

System:
    process_list:
    - 1
    - 2
    - 3
    output_path: ./plugin_processed
    output_rootname: neurone
    output_suffix: ''
Plugin_setup:
    pixel_size: 1.0
    input_mrc:
    - ./aligned/neurone_0001/neurone_0001.st
    - ./aligned/neurone_0002/neurone_0002.st
    - ./aligned/neurone_0003/neurone_0003.st
    output_mrc:
    - ./aligned/neurone_0001/neurone_0001_ali.mrc
    - ./aligned/neurone_0002/neurone_0002_ali.mrc
    - ./aligned/neurone_0003/neurone_0003_ali.mrc
    tilt_angles:
    - ./aligned/neurone_0001/neurone_0001.rawtlt
    - ./aligned/neurone_0002/neurone_0002.rawtlt
    - ./aligned/neurone_0003/neurone_0003.rawtlt
    choice: a

Here, the process_list, input_mrc, output_mrc, and tilt_angles fields are populated automatically in the plugin.update_yaml function in plugin.py.

To create the template plugin.yaml file, we add the following to params.py:

def new_plugin_yaml(args):
    """
    Subroutine to create yaml file for plugin

    ARGS:
    args (Namespace) :: Namespace containing user parameter inputs
    """

    plugin_yaml_name = args.project_name.value + '_plugin.yaml'

    plugin_yaml_dict = {
        'System': {
            'process_list': None,
            'output_path': str(args.output_path.value),
            'output_rootname': args.project_name.value if args.rootname.value is None else args.rootname.value,
            'output_suffix': args.suffix.value,
        },

        'Plugin_setup': {
            'pixel_size': args.pixel_size.value,
            'input_mrc': None,
            'output_mrc': None,
            'tilt_angles': None,
            'choice': args.choice.value,
        }
    }

    with open(plugin_yaml_name, 'w') as f:
        yaml.dump(plugin_yaml_dict, f, indent=4, sort_keys=False)

Don't worry about populating all the fields just now, just fill in those which can be obtained from the args, i.e., the user arguments captured from magicgui.

Next, we will need to create the plugin.py file, which will hold most of the logic needed for the plugin. We will go through this in two parts, here we will write the sections which deal with creating the yaml to setup the plugin.

Add the create_yaml and update_yaml methods to plugin.py.

def update_yaml(args):
    """Method to update yaml file

    Here we set the process list, specific filepaths, and check inputs

    ARGS:
    args (magicgui.FunctionGUI) :: magicgui object containing user input
    """
    # Read in template yaml file
    plugin_yaml_name = f"{args["project_name"]}_plugin.yaml"
    plugin_params = prmMod.read_yaml(
        project_name=args["project_name"],
        filename=plugin_yaml_name
    )

    # Set input filepaths, can be by searching for specific filenames in the input mrc folder, then update the yaml file
    input_files = glob(f"{args['input_mrc']}/*.mrc")
    plugin_params.params["Plugin_setup"]["input_mrc"] = input_files

    # Can do similar for the output filepaths

    # Set process list, this is usually the tilt series index, can be parsed from filepaths

    # Update yaml
    with open(Path(plugin_yaml_name), "w") as f:
        yaml.dump(plugin_params.params, f, indent=4, sort_keys=False)

def create_yaml(input_mgNS=None):
    """
    Subroutine to create new yaml file for Plugin
    """

    # Parse user inputs
    if input_mgNS is None:
        args = mgMod.get_args_plugin.show(run=True).asdict()
    else:
        args = input_mgNS

    # Create the yaml file, then automatically update it
    prmMod.new_plugin_yaml(args)
    update_yaml(args)

Now the functions to create the yaml should be in place, and we can move on to generating and running commands based on the yaml file.

Step 3: Functions to do plugin's work

The Plugin class in plugin.py creates the results folders, generates commands to run the plugin, runs the plugin, and saves the metadata.

Let's write the Plugin class in stages. First, set up the Plugin class. All Ot2Rec plugin classes have the same fundamental attributes

• project_name
• params_in: parameters read from the yaml file
• logger_in: a logger object to create the logfile o2r_plugin.log
• md_out: dictionary of metadata to pass to plugin_mdout.yaml.

In the __init__, we also want to set up the results folder structure, which is done in the _get_internal_metadata() function.

class Plugin:
    def __init__(self, project_name, params_in, logger_in):
        self.proj_name = project_name
        self.params = params_in.params
        self.logObj = logger_in
        self.md_out = {}

        self._get_internal_metadata()

    def _get_internal_metadata(self):
        """Prep internal metadata for processing and checking
        
        ** See other plugins for code to reuse **

        Generally, this sets the input and output filepaths and creates results folders.

        You can also pass information to the md_out, which is a metadata yaml file written after processing. e.g., putting in filepaths to the results.
        """
        for curr_ts in self.params["System"]["process_list"]:
            subfolder = (f"{self.basis_folder}/"
                         f"{self.rootname}_{curr_ts:04d}{self.suffix}")
            os.makedirs(subfolder, exist_ok=True)
        
            self.md_out["plugin_output_dir"][curr_ts] = subfolder
            self.md_out["plugin_output_file"][curr_ts] = f"{subfolder}/example.st

Next, we want to generate the commands needed to run the plugin. These are the same commands we would use if we were to use the plugin program directly (e.g., IMOD, AreTomo) in the terminal.

Add the following to the Plugin class.

    def _get_plugin_command(self, i):
        """Get command to call an external process for the i-th tilt series"""
        cmd = [
            "plugin",
            "-input",
            self.params["Plugin_setup"]["input_mrc"][i],
            "-output",
            self.params["Plugin_setup"]["output_mrc"][i],
            "-tiltangle",
            self.params["Plugin_setup"]["tilt_angles"][i],
            "-choice",
            self.params["Plugin_setup"]["choice]
        ]
        
        return cmd

We also want to add the runner functions to the Plugin class.

    def _run_plugin(self, i):
        """Run the plugin for the i-th tilt series"""
        cmd = self._get_plugin_command(i)
        plugin_run = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            encoding="ascii",
            check=True
        )
        self.logObj(plugin_run.stdout) # save stdout to log
    
    def run_plugin_all(self):
        for i, ts in enumerate(self.params["System"]["process_list"]):
            self._run_plugin(i)
        self.export_metadata()

    def export_metadata(self):
        yaml_file = self.proj_name + "_plugin_mdout.yaml"
        with open(yaml_file, "w") as f:
            yaml.dump(self.md_out, f, indent=4, sort_keys=False)

Lastly, we will add the run function to plugin.py, which is called by o2r.plugin.run.

def run():
    """
    Method to run plugin
    """
    # argparser to collect the project name
    parser = argparse.ArgumentParser()
    parser.add_argument("project_name",
                        type=str,
                        help="Name of current project")
    args = parser.parse_args()

    # Check if prerequisite files exist
    plugin_yaml_name = f"{args["project_name"]}_plugin.yaml"
    if not os.path.isfile(plugin_yaml_name):
        raise IOError("Error in Ot2Rec.main.run_plugin: plugin yaml file not found.")

    # Read in config and metadata
    plugin_config = prmMod.read_yaml(
        project_name=args.project_name,
        filename=plugin_yaml_name
    )

    # Create Logger object
    logger = logMod.Logger(log_path="o2r_plugin.log")

    # Create Plugin object
    plugin_obj = Plugin(
        project_name=args.project_name,
        params_in=plugin_config,
        logger_in=logger
    )

    # Run AreTomo commands
    plugin_obj.run_plugin_all()

Step 4: Add CLI binding

Lastly, add the plugin's entry points to setup.py to use it in the command line. In setup.py under entry_points, add:

"o2r.plugin.new=Ot2Rec.plugin:create_yaml",
"o2r.plugin.run=Ot2Rec.plugin:run",

Step 5: Share your plugin

If you'd like your plugin to be part of the public Ot2Rec, raise a pull request and we'll review your code and merge it if everything works well.