Source code for pdb2pqr.inputgen

"""Create an APBS input file using :mod:`psize` data.

.. codeauthor::  Todd Dolinsky
.. codeauthor::  Nathan Baker
"""
import pickle
import logging
import argparse
from pathlib import Path
from . import psize
from .config import TITLE_STR


_LOGGER = logging.getLogger(__name__)




class Elec:
    """An object for the ELEC section of an APBS input file."""



    def __init__(
        self, pqrpath, size, method, asyncflag, istrng=0, potdx=False
    ):
        """Initialize object.

        .. todo::  Remove hard-coded parameters.

        :param pqrpath:  path to PQR file
        :type pqrpath:  str
        :param size:  parameter sizing object
        :type size:  Psize
        :param method:  solution method (e.g., mg-para, mg-auto, etc.)
        :type method:  str
        :param asyncflag:  perform an asynchronous parallel focusing
            calculation
        :type asyncflag:  bool
        :param istrng:  ionic strength/concentration (M)
        :type istring:  float
        :param potdx:  whether to write out potential information in DX format
        :type potdx:  bool
        """
        # If this is an async or parallel calc, we want to use
        # the per-grid dime rather than the global dime.
        self.dime = size.ngrid
        gmem = (
            200.0
            * self.dime[0]
            * self.dime[1]
            * self.dime[2]
            / 1024.0
            / 1024.0
        )
        if method == "":  # method not named - use ceiling
            if gmem > size.gmemceil:
                method = "mg-para"
            else:
                method = "mg-auto"
        if method == "mg-para":
            self.dime = size.getSmallest()
        self.method = method
        self.istrng = istrng
        self.glen = size.coarse_length
        self.cglen = size.coarse_length
        self.fglen = size.fine_length
        self.pdime = size.proc_grid
        self.label = ""
        self.nlev = 4
        self.ofrac = 0.1
        self.async_ = 0
        self.asyncflag = asyncflag
        self.cgcent = "mol 1"
        self.fgcent = "mol 1"
        self.gcent = "mol 1"
        self.mol = 1
        self.lpbe = 1
        self.npbe = 0
        self.bcfl = "sdh"
        # TODO - where did these very arbitrary numbers come from?
        self.ion = [[-1, 1.815], [1, 1.875]]  # Multiple ions possible
        self.pdie = 2.0
        self.sdie = 78.54
        self.srfm = "smol"
        self.chgm = "spl2"
        self.sdens = 10.0
        self.srad = 1.4
        self.swin = 0.3
        self.temp = 298.15
        self.gamma = 0.105
        self.calcenergy = "total"
        self.calcforce = "no"
        if potdx:
            self.write = [["pot", "dx", pqrpath]]
        else:
            # Multiple write statements possible
            self.write = [["pot", "dx", "pot"]]


    def __str__(self):
        text = f"elec {self.label}\n"
        text += f"    {self.method}\n"
        text += f"    dime {int(self.dime[0])} {int(self.dime[1])} "
        text += f"{int(self.dime[2])}\n"
        if self.method == "mg-auto":
            text += f"    cglen {self.cglen[0]:.4f} {self.cglen[1]:.4f} "
            text += f"{self.cglen[2]:.4f}\n"
            text += f"    fglen {self.fglen[0]:.4f} {self.fglen[1]:.4f} "
            text += f"{self.fglen[2]:.4f}\n"
            text += f"    cgcent {self.cgcent}\n"
            text += f"    fgcent {self.fgcent}\n"
        elif self.method == "mg-manual":
            text += f"    glen {self.glen[0]:.3f} {self.glen[1]:.3f} "
            text += f"{self.glen[2]:.3f}\n"
            text += f"    gcent {self.gcent}\n"
        elif self.method == "mg-para":
            text += f"    pdime {int(self.pdime[0])} {int(self.pdime[1])} "
            text += f"{int(self.pdime[2])}\n"
            text += f"    ofrac {self.ofrac:.1f}\n"
            text += f"    cglen {self.cglen[0]:.4f} {self.cglen[1]:.4f} "
            text += f"{self.cglen[2]:.4f}\n"
            text += f"    fglen {self.fglen[0]:.4f} {self.fglen[1]:.4f} "
            text += f"{self.fglen[2]:.4f}\n"
            text += f"    cgcent {self.cgcent}\n"
            text += f"    fgcent {self.fgcent}\n"
            if self.asyncflag:
                text += f"    async {self.async_}\n"
        text += f"    mol {int(self.mol)}\n"
        text += "    lpbe\n" if self.lpbe else "    npbe\n"
        text += f"    bcfl {self.bcfl}\n"
        if self.istrng > 0:
            for ion in self.ion:
                text += f"    ion charge {ion[0]:.2f} conc {self.istrng:.3f} "
                text += f"radius {ion[1]:.4f}\n"
        text += f"    pdie {self.pdie:.4f}\n"
        text += f"    sdie {self.sdie:.4f}\n"
        text += f"    srfm {self.srfm}\n"
        text += f"    chgm {self.chgm}\n"
        text += f"    sdens {self.sdens:.2f}\n"
        text += f"    srad {self.srad:.2f}\n"
        text += f"    swin {self.swin:.2f}\n"
        text += f"    temp {self.temp:.2f}\n"
        text += f"    calcenergy {self.calcenergy}\n"
        text += f"    calcforce {self.calcforce}\n"
        for write in self.write:
            text += f"    write {write[0]} {write[1]} {write[2]}\n"
        text += "end\n"
        return text





class Input:
    """Each object of this class is one APBS input file."""



    def __init__(
        self, pqrpath, size, method, asyncflag, istrng=0, potdx=False
    ):
        """Initialize the input file class.

        Each input file contains a PQR name, a list of elec objects, and a
        list of strings containing print statements.
        For starters, assume two ELEC statements are needed, one for the
        inhomgenous and the other for the homogenous dielectric calculations.

        .. note::
            This assumes you have already run psize, either by
            :func:`size.run_psize(...)` or :func:`size.parse_string(...)`
            followed by :func:`size.set_all()`.

        :param pqrpath:  path to PQR file
        :type pqrpath:  str
        :param size:  parameter sizing object
        :type size:  Psize
        :param method:  solution method (e.g., mg-para, mg-auto, etc.)
        :type method:  str
        :param asyncflag:  perform an asynchronous parallel focusing
            calculation
        :type asyncflag:  bool
        :param istrng:  ionic strength/concentration (M)
        :type istring:  float
        :param potdx:  whether to write out potential information in DX format
        :type potdx:  bool
        """
        self.pqrpath = Path(pqrpath)
        self.pqrname = self.pqrpath.name
        self.asyncflag = asyncflag
        # Initialize variables to default elec values
        elec1 = Elec(pqrpath, size, method, asyncflag, istrng, potdx)
        if not potdx:
            elec2 = Elec(pqrpath, size, method, asyncflag, istrng, potdx)
            setattr(elec2, "sdie", 2.0)
            setattr(elec2, "write", [])
        else:
            elec2 = ""
        self.elecs = [elec1, elec2]
        if not potdx:
            self.prints = ["print elecEnergy 2 - 1 end"]
        else:
            self.prints = ["print elecEnergy 1 end"]


    def __str__(self):
        text = "read\n"
        text += f"    mol pqr {self.pqrname}\n"
        text += "end\n"
        for elec in self.elecs:
            text += str(elec)
        for prints in self.prints:
            text += prints
        text += "\nquit\n"
        return text



    def print_input_files(self, output_path):
        """Generate the input file(s) associated with this object.

        :param output_path:  location for generated files
        :type output_path:  str
        """
        path = Path(output_path)
        base_name = path.stem
        if self.asyncflag:
            outname = base_name + "-para.in"
            # Temporarily disable async flag
            for elec in self.elecs:
                elec.asyncflag = False
            with open(outname, "wt") as out_file:
                out_file.write(str(self))
            # Now make the async files
            elec = self.elecs[0]
            nproc = elec.pdime[0] * elec.pdime[1] * elec.pdime[2]
            for i in range(int(nproc)):
                outname = base_name + f"-PE{i}.in"
                for elec in self.elecs:
                    elec.asyncflag = True
                    elec.async_ = i
                with open(outname, "wt") as out_file:
                    out_file.write(str(self))
        else:
            with open(path, "wt") as out_file:
                out_file.write(str(self))




    def dump_pickle(self):
        """Make a Python pickle associated with the APBS input parameters.

        .. todo::  is this function still useful?
        """
        base_pqr_name = self.pqrpath.stem
        outname = base_pqr_name + "-input.p"
        with open(outname, "wb") as pfile:
            pickle.dump(self, pfile)






def split_input(filename):
    """Split the parallel input file into multiple async file names.

    :param filename:  the path to the original parallel input file
    :type filename:  str
    """
    nproc = 0
    with open(filename, "rt") as file_:
        text = ""
        while True:
            line = file_.readline()
            if not line:
                break
            text += line
            line = line.strip()
            if line.startswith("pdime"):  # Get # Procs
                words = line.split()
                nproc = int(words[1]) * int(words[2]) * int(words[3])
    if nproc == 0:
        errstr = f"{filename} is not a valid APBS parallel input file!\n"
        errstr = errstr + (
            "The inputgen script was unable to asynchronize this file!"
        )
        raise RuntimeError(errstr)
    base_pqr_name = Path(filename).stem
    for iproc in int(range(nproc)):
        outname = base_pqr_name + f"-PE{iproc}.in"
        outtext = text.replace("mg-para\n", f"mg-para\n    async {iproc}\n")
        with open(outname, "w") as outfile:
            outfile.write(outtext)





def build_parser():
    """Build argument parser."""
    desc = f"{TITLE_STR}\ninputgen: generating APBS input files since "
    desc += "(at least) 2004"
    parse = argparse.ArgumentParser(
        description=desc,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parse.add_argument(
        "--asynch",
        action="store_true",
        help="perform an asynchronous parallel calculation.",
    )
    parse.add_argument(
        "--split",
        action="store_true",
        help=(
            "split an existing parallel input file to multiple "
            "async input files."
        ),
    )
    parse.add_argument(
        "--potdx",
        action="store_true",
        help=("create an input to compute an electrostatic potential map."),
    )
    parse.add_argument(
        "--method",
        help=("force output file to write a specific APBS ELEC method."),
        choices=["para", "auto", "manual", "async"],
    )
    parse.add_argument(
        "--cfac",
        type=float,
        default=psize.CFAC,
        help=(
            "factor by which to expand molecular dimensions to "
            "get coarse grid dimensions."
        ),
    )
    parse.add_argument(
        "--fadd",
        type=float,
        default=psize.FADD,
        help=(
            "amount to add to molecular dimensions to get fine "
            "grid dimensions."
        ),
    )
    parse.add_argument(
        "--space",
        type=float,
        default=psize.SPACE,
        help="desired fine mesh resolution",
    )
    parse.add_argument(
        "--gmemfac",
        type=int,
        default=psize.GMEMFAC,
        help=(
            "number of bytes per grid point required for sequential "
            "MG calculation"
        ),
    )
    parse.add_argument(
        "--gmemceil",
        type=int,
        default=psize.GMEMCEIL,
        help=(
            "max MB allowed for sequential MG calculation; adjust "
            "this to force the script to perform faster calculations "
            "(which require more parallelism)"
        ),
    )
    parse.add_argument(
        "--ofrac",
        type=float,
        default=psize.OFRAC,
        help="overlap factor between mesh partitions (parallel)",
    )
    parse.add_argument(
        "--redfac",
        type=float,
        default=psize.REDFAC,
        help=(
            "the maximum factor by which a domain dimension can "
            "be reduced during focusing"
        ),
    )
    parse.add_argument(
        "--istrng", help="Ionic strength (M); Na+ and Cl- ions will be used"
    )
    parse.add_argument("filename")
    return parse





def main():
    """Main driver"""
    parser = build_parser()
    args = parser.parse_args()
    size = psize.Psize()
    filename = args.filename
    if args.split:
        split_input(filename)
    else:
        size.run_psize(filename)
        input_ = Input(
            filename, size, args.method, args.asynch, args.istrng, args.potdx
        )
        path = Path(filename)
        output_path = path.parent + path.stem + Path(".in")
        input_.print_input_files(output_path)



if __name__ == "__main__":
    main()