Implementing a VVUQ campaign into the FabNESO plugin

FabNESO/__init__.py

@@ -1,9 +1,23 @@
 """Neptune Exploratory SOftware (NESO) plugin for FabSim3."""
 
 try:
-    from .tasks import neso, neso_ensemble, neso_vbmc, neso_write_field
+    from .tasks import (
+        neso,
+        neso_ensemble,
+        neso_vbmc,
+        neso_vvuq_campaign,
+        neso_write_field,
+        single_run_vvuq,
+    )
 
-    __all__ = ["neso", "neso_ensemble", "neso_vbmc", "neso_write_field"]
+    __all__ = [
+        "neso",
+        "neso_ensemble",
+        "neso_vbmc",
+        "neso_write_field",
+        "neso_vvuq_campaign",
+        "single_run_vvuq",
+    ]
 
 except ImportError:
     import warnings

FabNESO/run_vvuq_instance.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+
+echo fabsim $1 single_run_vvuq:$PWD,mesh_file=$3,solver=$4,processes=$5,nodes=$6,cpus_per_process=$7,wall_time=$8,neso_outfile=$9
+
+fabsim $1 single_run_vvuq:$PWD,mesh_file=$3,solver=$4,processes=$5,nodes=$6,cpus_per_process=$7,wall_time=$8,neso_outfile=$9

FabNESO/tasks.py

@@ -7,11 +7,14 @@
 import re
 import shutil
 import time
+from ast import literal_eval
 from contextlib import nullcontext
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Any
 
+import chaospy as cp
+import easyvvuq as uq
 import numpy as np
 import pyvbmc
 
@@ -579,3 +582,199 @@ def log_density(
         (config_dict["initial_run"] - observed_results["field_value"]) ** 2
         / (2 * config_dict["observation_noise_std"] ** 2)
     ).sum()
+
+
+def _parse_vvuq_parameters(
+    vvuq_param_file: Path,
+) -> tuple[dict[Any, dict[str, float]], dict[str, Any], list[str]]:
+    # read the inpurt file
+    with Path.open(vvuq_param_file) as f:
+        data = f.read()
+    d = literal_eval(data)
+    vvuq_params = {}
+    parameter_variations = {}
+    for parameter, param_values in d["parameters"].items():
+        vvuq_params[parameter] = {
+            "type": "float",
+            "default": param_values[0],
+        }
+        vary_length = 3
+        if len(param_values) == vary_length:
+            # Make the chaospy distribution. Possibly to be able to edit this later?
+            parameter_variations[parameter] = cp.Uniform(
+                param_values[1], param_values[2]
+            )
+    output_columns = d["output_columns"]
+    return (vvuq_params, parameter_variations, output_columns)
+
+
+@fab.task
+@fab.load_plugin_env_vars("FabNESO")
+def single_run_vvuq(
+    config: str,
+    mesh_file: str = "mesh.xml",
+    solver: str = "Electrostatic2D3V",
+    processes: str | int = 4,
+    nodes: str | int = 1,
+    cpus_per_process: str | int = 1,
+    wall_time: str = "00:15:00",
+    neso_outfile: str = "Electrostatic2D3V_line_field_evaluations.h5part",
+) -> None:
+    """
+    Run a single instance of the VVUQ campaign.
+
+    Run an instance of the VVUQ campaign using configurations in the designated solver.
+    The conditions and mesh are, by design, the defaults. When the iteration is done,
+    copy the NESO output to the local directory for later decoding through VVUQ.
+
+
+    Args:
+        config: Directory with single run configuration information.
+
+    Keyword Args:
+        solver: Which NESO solver to use.
+        mesh_file_name: Name of mesh XML in configuration directory.
+        processes: Number of processes to run.
+        nodes: Number of nodes to run on. Only applicable when running on a multi-node
+            system.
+        cpus_per_process: Number of processing units to use per process. Only
+            applicable when running on a multi-node system.
+        wall_time: Maximum time to allow job to run for. Only applicable when submitting
+            to a job scheduler.
+        neso_outfile: Name of the output file the NESO solver produces, to copy across
+            to the local directory.
+
+    """
+    # If we're running remotely, tell the submission to wait until done
+    if "archer2" in fab.env.remote:
+        fab.update_environment(
+            {
+                "job_dispatch": "cd /work/$project/$project/$username ; sbatch --wait",
+            }
+        )
+
+    # Make a temporary directory to store the configs in - they can be deleted when the
+    # run is done
+    with TemporaryDirectory(dir=config) as temporary_config_dir:
+        config_path = Path(temporary_config_dir)
+
+        # Copy the mesh and conditions files here for FabSIM to read it properly
+        shutil.copyfile(mesh_file, config_path / Path(mesh_file).name)
+        shutil.copyfile("conditions.xml", config_path / "conditions.xml")
+
+        # Run an instance of NESO
+        neso(
+            temporary_config_dir,
+            solver=solver,
+            mesh_file_name=Path(mesh_file).name,
+            processes=processes,
+            nodes=nodes,
+            cpus_per_process=cpus_per_process,
+            wall_time=wall_time,
+        )
+
+    # Retrieve the results
+    fab.fetch_results()
+
+    # Move the results to the local VVUQ directory for further processing
+    local_results_dir = Path(fab.env.job_results_local) / template(
+        fab.env.job_name_template
+    )
+    shutil.copyfile(local_results_dir / neso_outfile, neso_outfile)
+
+
+@fab.task
+@fab.load_plugin_env_vars("FabNESO")
+def neso_vvuq_campaign(  # noqa: PLR0913
+    vvuq_config: str,
+    conditions_template: str = "conditions.template",
+    vvuq_parameters: str = "vvuq_parameters.txt",
+    vvuq_script: str = "FabNESO/run_vvuq_instance.sh",
+    neso_mesh_file: str = "mesh.xml",
+    output_dir: str = "neso_pce/",
+    neso_outfile: str = "Electrostatic2D3V_line_field_evaluations.h5part",
+    solver: str = "Electrostatic2D3V",
+    processes: str | int = 4,
+    nodes: str | int = 1,
+    cpus_per_process: str | int = 1,
+    wall_time: str = "00:15:00",
+    *,
+    run_vvuq_sequential: bool = True,
+) -> None:
+    """
+    Run a VVUQ campaign on a NESO solver.
+
+    Args:
+        vvuq_config: Directory with VVUQ parameters, conditions template file, and mesh
+            configuration.
+
+    Keyword Args:
+        conditions_template: A template NESO conditions file. Parameters that are to be
+            varied are indicated with a $ sign.
+        vvuq_parameters: A file in the vvuq_config directory containing a dict with
+            "parameters" that correspond to the parameters to replace in template file,
+            "output_columns" that are the NESO outputs that VVUQ should read for its
+            analysis
+        neso_mesh_file: Name of mesh XML in vvuq_config directory.
+        output_dir: Directory in which to carry out the VVUQ campaign.
+        neso_outfile: Name of the output file the NESO solver produces, to copy across
+            to the local directory.
+        solver: Which NESO solver to use.
+        processes: Number of processes to run.
+        nodes: Number of nodes to run on. Only applicable when running on a multi-node
+            system.
+        cpus_per_process: Number of processing units to use per process. Only
+            applicable when running on a multi-node system.
+        wall_time: Maximum time to allow job to run for. Only applicable when submitting
+            to a job scheduler.
+        run_vvuq_sequential: Set to False to parallelise. Default is True because this
+            can cause issues if unchecked.
+
+    """
+    path_to_vvuq_configs = Path(fab.find_config_file_path(vvuq_config))
+
+    # Extract the parameters for the conditions template and variations from the config
+    params, vary, outputs_to_read = _parse_vvuq_parameters(
+        path_to_vvuq_configs / vvuq_parameters
+    )
+
+    # Make an encoder
+    encoder = uq.encoders.GenericEncoder(
+        path_to_vvuq_configs / conditions_template,
+        delimiter="$",
+        target_filename="conditions.xml",
+    )
+
+    # Make the VVUQ decoder - this should be more generic!
+    decoder = uq.decoders.HDF5(
+        target_filename=neso_outfile,
+        output_columns=outputs_to_read,
+    )
+
+    # Get the path for this plugin,
+    vvuq_run_path = Path(fab.get_plugin_path("FabNESO")) / vvuq_script
+
+    # The script that runs the individual NESO instances
+    execute_local = uq.actions.ExecuteLocal(
+        f"{vvuq_run_path} {fab.env.remote} {vvuq_config} "
+        f"{path_to_vvuq_configs / neso_mesh_file} {solver} {processes} "
+        f"{nodes} {cpus_per_process} {wall_time} {neso_outfile}"
+    )
+
+    # Put together the actions required for the VVUQ campaign
+    actions = uq.actions.Actions(
+        uq.actions.CreateRunDirectory("/tmp", flatten=True),  # noqa: S108
+        uq.actions.Encode(encoder),
+        execute_local,
+        uq.actions.Decode(decoder),
+    )
+
+    # Make the VVUQ campaign
+    campaign = uq.Campaign(name=output_dir, params=params, actions=actions)
+
+    # Make and attach a sampler to the campaign
+    my_sampler = uq.sampling.PCESampler(vary)
+    campaign.set_sampler(my_sampler)
+
+    # Run VVUQ
+    campaign.execute(sequential=run_vvuq_sequential).collate()

config_files/vvuq_example/conditions.template

@@ -0,0 +1,77 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<NEKTAR xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:noNamespaceSchemaLocation="http://www.nektar.info/schema/nektar.xsd">
+
+    <EXPANSIONS>
+        <E COMPOSITE="C[1]" NUMMODES="2" TYPE="MODIFIED" FIELDS="u" />
+        <E COMPOSITE="C[1]" NUMMODES="2" TYPE="MODIFIED" FIELDS="rho" />
+    </EXPANSIONS>
+
+    <CONDITIONS>
+
+        <SOLVERINFO>
+            <I PROPERTY="EQTYPE" VALUE="PoissonPIC" />
+            <I PROPERTY="Projection" VALUE="Continuous" />
+        </SOLVERINFO>
+
+        <GLOBALSYSSOLNINFO>
+            <V VAR="u">
+            <I PROPERTY="GlobalSysSoln" VALUE="IterativeStaticCond" />
+            <I PROPERTY="IterativeSolverTolerance" VALUE="1e-8"/>
+            </V>
+        </GLOBALSYSSOLNINFO>
+
+        <PARAMETERS>
+            <P> Lambda = 0.0 </P>
+            <P> epsilon = -1.0 </P>
+            <P> num_particles_total = $NUMBER_PARTICLES </P>
+            <P> num_particles_per_cell = -1 </P>
+            <P> particle_time_step = 0.001 </P>
+            <P> particle_num_time_steps = $TIME_STEPS </P>
+            <P> particle_num_write_particle_steps = 0 </P>
+            <P> particle_num_write_field_energy_steps = 20 </P>
+            <P> particle_num_write_field_steps = $WRITE_FIELD_STEPS </P>
+            <P> particle_num_print_steps = 40 </P>
+            <P> particle_distribution_position = 2 </P>
+            <P> particle_initial_velocity = $INIT_VELOCITY </P>
+            <P> particle_charge_density = $CHARGE_DENSITY </P>
+            <P> particle_number_density = $NUMBER_DENSITY </P>
+            <P> line_field_deriv_evaluations_step = $LINE_STEP </P>
+            <P> line_field_deriv_evaluations_numx = $LINE_NUMX </P>
+            <P> line_field_deriv_evaluations_numy = $LINE_NUMY </P>
+        </PARAMETERS>
+
+        <VARIABLES>
+            <V ID="0"> u </V>
+            <V ID="1"> rho </V>
+        </VARIABLES>
+
+        <BOUNDARYREGIONS>
+            <B ID="1"> C[100] </B>
+            <B ID="2"> C[200]  </B>
+            <B ID="3"> C[300] </B>
+            <B ID="4"> C[400] </B>
+        </BOUNDARYREGIONS>
+
+        <BOUNDARYCONDITIONS>
+            <REGION REF="1">
+                <P VAR="u" VALUE="[3]" />
+                <P VAR="rho" VALUE="[3]" />
+            </REGION>
+            <REGION REF="2">
+                <P VAR="u" VALUE="[4]" />
+                <P VAR="rho" VALUE="[4]" />
+            </REGION>
+            <REGION REF="3">
+                <P VAR="u" VALUE="[1]" />
+                <P VAR="rho" VALUE="[1]" />
+            </REGION>
+            <REGION REF="4">
+                <P VAR="u" VALUE="[2]" />
+                <P VAR="rho" VALUE="[2]" />
+            </REGION>
+        </BOUNDARYCONDITIONS>
+
+    </CONDITIONS>
+
+</NEKTAR>