Allow the BSS.FreeEnergy.Relative to run Gromacs process with a single MPI rank

[xiki-tempula]

Is your feature request related to a problem? Please describe.
I'm interested in changing something in Gromacs mdp generation and wish to test it with a simple system.
I bumped into the #305

Describe the solution you'd like
I wonder if I could do something to allow each window of BSS.FreeEnergy.Relative to be run with a single MPI rank
with interface

        # Benzene.
        m0 = BSS.Parameters.openff_unconstrained_2_0_0(
            "c1ccccc1").getMolecule()
        # Toluene.
        m1 = BSS.Parameters.openff_unconstrained_2_0_0(
            "Cc1ccccc1").getMolecule()
        atom_mapping = BSS.Align.matchAtoms(m0, m1)
        m0 = BSS.Align.rmsdAlign(m0, m1, atom_mapping)
        merged = BSS.Align.merge(m0, m1)
        solvated = BSS.Solvent.tip3p(molecule=merged,
                                     box=3 * [3 * BSS.Units.Length.nanometer])
        protocol = BSS.Protocol.FreeEnergy(lam=0.0,
                 lam_vals=None,
                 min_lam=0.0,
                 max_lam=1.0,
                 num_lam=11,
                 timestep=_Types.Time(2, "femtosecond"),
                 runtime=_Types.Time(0, "femtosecond"),
                 temperature_start=_Types.Temperature(300, "kelvin"),
                 temperature_end=_Types.Temperature(300, "kelvin"),
                 temperature=None,
                 pressure=None,
                 report_interval=200,
                 restart_interval=1000,
                 restraint=None,
                 force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,
                 restart=False,
                 perturbation_type="full")
        freenrg = BSS.FreeEnergy.Relative(system, protocol, engine='GROMACS', )
        freenrg.setArg({"-ntmpi": 1})
        freenrg.run()
        freenrg.wait()

[lohedges]

Hello there,

There currently isn't an exposed way to modify the run arguments as you can for a Process object. However, there is a private _uddate_run_args method that can be used to do the same thing. Since this updates the entire dictionary of arguments for all processes you'd need to do something like the following.

...

freenrg = BSS.FreeEnergy.Relative(system, protocol, engine="GROMACS")

# Get the args to "gmx mdrun". Here we assume that they are all the same.
args = freenrg._runner.processes()[0].getArgs()

# Use a single MPI thread/rank.
args["-ntmpi"] = 1

# Set the arguments for all processes.
freenrg._update_run_args(args)

Since freenrg.run() will run all lambda windows sequentially by default, I wouldn't really use this for testing. If you wanted a short unit test, then I'd just use Process.FreeEnergy directly to run a single lambda window as a test. If you use this approach, then you can just use setArg on the Process directly.

(Be careful restricting threads with GROMACS since the behaviour, e.g. whether it runs, depends on how GROMACS was compiled and on what type of system you are running. I'm not sure the specific thread requirement for FEP with GROMACS (OMP or MPI threads) but I know, for example, that there can be issues restricting the number of threads for metadynamics simulations.)

[xiki-tempula]

@lohedges Thank you for the reply. I wonder if there is a plan of adding one?
As for gromacs, one might wish to have the flag -update gpu to use GPU update or -replex 1000 -nex 1000000 to run replica exchange, so it would be helpful, if one could pass such an argument.

[lohedges]

Yes, I hope to make things more flexible in future. As I say, in general it's hard to know what options would be suitable for a particular build or environment. I plan on adding a sub-package where a user can set information related to their hardware resources, scheduler, etc, to make automatic setting of appropriate arguments easier.

I also aim to add in @msuruzhon's optional config and extra arguments approach.

[xiki-tempula]

@lohedges Thanks for the comment. I wonder if it is possible to have an interface similar to

        freenrg = BSS.FreeEnergy.Relative(system, protocol, engine='GROMACS', )
        freenrg.setArg({"-ntmpi": 1})

Where the Arg provided here will simply be appended to the existing freenrg._runner.processes()[0].getArgs() so the user could set the args manually?

[lohedges]

Yes, that would be a good approach at first. I'll aim to mirror the functionality from BioSimSpace.Process so that the user can edit the config file and arguments using the same interface. We originally didn't add this at this level since we weren't sure whether we wanted to allow the user to specify the engine, which we now do.

[xiki-tempula]

@lohedges Thank you for the explanation.