goPEST 0.0.11__py3-none-any.whl

gopest/pest_model.py

@@ -0,0 +1,245 @@
+""" This is the actual code that wraps TOUGH2 model run with pre- and post-
+processing, so all PEST's model/batch file call this.
+
+Use with optiona flags:
+    python pest_model.py [--svda] [--obsreref]
+
+Run with "--svda" flag with additional command of parcalc for SVDassist runs.
+Flag "--obsreref" is for observation re-referencing which reset save to incon
+and overwrites the master's incon, which will be used by all subsequent model
+runs.  Flag "--test-update" will cause each slave to save a unique
+real_model.save file back to master directory.  This allows "--obsreref" to
+later select the best save/incon to start with.
+
+To generate/overwrite/fix the model/batch files, use:
+    python make_batch_files.py
+"""
+
+import time
+import sys
+import glob
+import os
+from os import devnull, system, remove, sep, path
+from shutil import copy2
+from shutil import Error
+from time import sleep
+
+from numpy.testing import assert_approx_equal
+
+from gopest.run_ns_pr import run_ns_pr
+from gopest.par import generate_real_model
+from gopest.obs import read_from_real_model
+
+from gopest.common import config
+from gopest.common import runtime
+
+def get_master_dir():
+    # if config['mode'] != 'local':
+    with open('_master_dir', 'r') as f:
+        line = f.readlines()[0].strip()
+        return line
+
+def get_pest_dir():
+    with open('_pest_dir', 'r') as f:
+        line = f.readlines()[0].strip()
+        return line
+
+def get_t2():
+    with open('_tough2', 'r') as f:
+        line = f.readlines()[0].strip()
+        return line
+
+def get_slave_id():
+    try:
+        with open('_procid', 'r') as f:
+            line = f.readlines()[0].strip()
+            return line
+    except:
+        return '0'
+
+def par_match(pf1, pf2):
+    matched = False
+    with open(pf1,'r') as a:
+        with open(pf2,'r') as b:
+            # once open successfully, assume equal, until something fails
+            matched = True
+            try:
+                for aa,bb in zip(a,b):
+                    ax, bx = float(aa.split(',')[0]), float(bb.split(',')[0])
+                    assert_approx_equal(ax, bx, significant=7)
+            except AssertionError:
+                matched = False
+    return matched
+
+def main(obsreref, svda, testup, local, skiprun, useobf, sendbad, skippr, hdf5, waiwera):
+    fgeo = runtime['filename']['geom']
+    fsave = runtime['filename']['save']
+    fincon = runtime['filename']['incon']
+    fdato = runtime['filename']['dat_orig']
+    fdats = runtime['filename']['dat_seq']
+    flsts = runtime['filename']['lst_seq']
+
+    print("----- Running " + " ".join(sys.argv[1:]))
+    if local:
+        master_dir = '.'
+    else:
+        master_dir = get_master_dir()
+    print("  --- Clean up pest_model.obf")
+    if path.isfile('pest_model.obf'):
+        remove('pest_model.obf')
+
+    ### skips everything if use obf directly
+    if useobf:
+        raise Exception('Should use pest_hp file distribution 17/12/2022')
+        if path.isfile('pest_model.obf.use'):
+            copy2(master_dir + sep + 'pest_model.obf.use', 'pest_model.obf')
+            print("Found pest_model.obf.use , skips everything.")
+            return
+        else:
+            print("Error, cannot find existing pest_model.obf.use file")
+
+    ### SVD-assist
+    if svda:
+        print("  --- PARCALC")
+        try:
+            remove('pest_model.dat')
+        except:
+            print("pest_model.dat probably does not exist")
+        PARCALC = path.join(get_pest_dir(), 'parcalc')
+        system(PARCALC + ' > ' + devnull)
+
+    ### goPESTpar
+    if not skiprun:
+        print("  --- goPESTpar")
+        generate_real_model(fdato, 'pest_model.dat', fdats[0])
+        # sleep(30)  # just in case shared file system slow
+
+    if obsreref:
+        if path.isfile('pest_model.obf'):
+            remove('pest_model.obf')
+        # get matching incon, if exist
+        for parf in glob.glob(master_dir + sep + 'pest_model.dat.*'):
+            if par_match(parf, 'pest_model.dat'):
+                matchname = path.splitext(parf)[1]
+                print("  --- found matched incon/pars %s from master dir, overwrite Master INCON" % matchname)
+                copy2(master_dir + sep + fincon + matchname, master_dir + sep + fincon)
+                copy2(master_dir + sep + 'pest_model.obf' + matchname, 'pest_model.obf')
+                break
+        # print("  --- remove all pairs from labmda tests after searching")
+        # for f in glob.glob(master_dir + sep + 'pest_model.obf.*'):
+        #     remove(f)
+        # for f in glob.glob(master_dir + sep + 'pest_model.dat.*'):
+        #     remove(f)
+        # for f in glob.glob(master_dir + sep + 'real_model.incon.*'):
+        #     remove(f)
+        if path.isfile('pest_model.obf'):
+            print("  --- use obf, skip actual model run")
+            return
+        else:
+            print("  --- could not find matching pars, obsreref continue with normal run")
+
+    ### RUN TOUGH2 model
+    if skiprun:
+        print("  --- skip actual TOUGH2 run")
+    else:
+        if not local:
+            print("  --- use master INCON")
+            try:
+                copy2(master_dir + sep + fincon, fincon)
+            except Error as e:
+                # OK if src and dst are the same file, simply skip.
+                print(e)
+
+        START_TIME = time.time()
+        print("  --- run_ns_pr()")
+        runok = run_ns_pr()
+        if obsreref:
+            if not local:
+                print("  --- reset Master INCON")
+                copy2('real_model.incon', master_dir + sep + 'real_model.incon')
+            else:
+                print("  --- .save file written as .incon")
+        if sendbad:
+            for f in glob.glob('bad_model_*'):
+                copy2(f, master_dir + sep + 'bad_model_slave' + get_slave_id() + '_' + f)
+        if not runok:
+            print("  --- run_ns_pr failed, skip goPESTobs, no obf, make sure lamforgive/derforgive is used.")
+            return
+        print('  --- run_ns_pr() complete after', (time.time() - START_TIME), 'seconds')
+
+    ### goPESTobs
+    # sleep(30)  # just in case shared file system slow
+    print("  --- goPESTobs")
+    read_from_real_model(fgeo, fdats[-1], flsts[-1], 'pest_model.obf', waiwera=waiwera)
+
+    if testup:
+        print("  --- store lambda test (save,obf,pars) pair:" + get_slave_id())
+        copy2(fsave, master_dir + sep + fincon + '.' + get_slave_id())
+        copy2('pest_model.dat', master_dir + sep + 'pest_model.dat.' + get_slave_id())
+        copy2('pest_model.obf', master_dir + sep + 'pest_model.obf.' + get_slave_id())
+
+def main_cli(argv=[]):
+    """ the main purpose of run-pest-model is to be called by PEST.
+
+    PEST needs:
+    - pest_model.tpl
+    - pest_model.ins
+
+    When PEST wants a forward run, PEST
+    1. uses pest_model.tpl to generate: pest_model.dat
+    2. runs the command: gopest run-pest-model
+    3. uses pest_model.ins to read data from: pest_model.obf
+
+    under the hood, run-pest-model's job is to:
+    a. get parameters from pest_model.dat to construct real_model.dat
+    b. run actual simulation, which generates real_model.h5
+    c. converts real_model.h5 into pest_model.obf
+
+    a. is done by goPESTpar
+    b. is done by run-forward
+    c. is done by goPESTobs
+    """
+    obsreref = False
+    svda = False
+    testup = False
+    skiprun = False
+    useobf = False
+    sendbad = True
+    skippr = False
+    waiwera = False
+    hdf5 = False
+
+    if config['simulator']['output-type'] == 'h5':
+        hdf5 = True
+    else:
+        hdf5 = False
+    
+    if config['simulator']['input-type'] == 'waiwera':
+        waiwera = True
+        hdf5 = True
+    else:
+        waiwera = False
+
+    # probably should be invoked from within gopest?
+    # local = config['pest']['mode'] == 'local'
+    local = False
+
+    skiprun = config['model']['skip']
+    skippr = config['model']['skip-pr']
+
+    print('pest_model.py running at ', os.getcwd())
+    
+    if len(argv) > 1:
+        if '--test-update' in argv[1:]:
+            testup = True
+        if '--obsreref' in argv[1:]:
+            obsreref = True
+        if '--svda' in argv[1:]:
+            svda = True
+        if '--use-obf' in argv[1:]:
+            # requires existing pest_model.obf.use (PEST will remove
+            # pest_model.obf, so use different name)
+            useobf = True
+        if '--local' in argv[1:]:
+            local = True
+    main(obsreref, svda, testup, local, skiprun, useobf, sendbad, skippr, hdf5, waiwera)

gopest/rename_latest_files.py

@@ -0,0 +1,35 @@
+import glob
+import os
+from shutil import copy2
+
+from gopest.common import config
+
+def current_iteration():
+    its = [int(f.split('.')[-1]) for f in glob.glob('%s.jco.*' % config['pest']['case-name'])]
+    its.append(0)
+    return max(its)
+
+def rename_latest_files(argv=[]):
+    """ PEST_HP's file distribution will copy files from the best update slave
+    to here.  They will be named as '*.999'.  This script renames these to the
+    latest iteration number.
+    """
+    print('gopest save-iter-model (rename_latest_files.py):')
+
+    # try:
+    #     copy2('real_model.incon.999', 'real_model.incon')
+    #     print("    real_model.incon.999 -> real_model.incon (copy)")
+    # except Exception as e:
+    #     print(" ", e)
+    #     print("  failed to copy incon file for next iteration!")
+
+    ii = current_iteration()
+    print('  Current iteration is %i' % ii)
+
+    for f in glob.glob('*.999'):
+        newname = f.replace('.999', '.%i' % ii)
+        print('    %s -> %s' % (f, newname))
+        try:
+            os.rename(f, newname)
+        except Exception as e:
+            print(e)

gopest/run_beopest.py

@@ -0,0 +1,205 @@
+import os
+import sys
+import shutil
+import time
+from subprocess import Popen, DEVNULL
+from multiprocessing import cpu_count
+
+from gopest.common import config as cfg
+from gopest.common import runtime
+
+NUM_SLAVES = cfg['pest']['num_slaves']
+SILENT_SLAVES = cfg['pest']['silent_slaves']
+PST_NAME = cfg['pest']['case-name']
+PORT = cfg['pest']['port']
+SWITCHES = " ".join(cfg['pest']['switches'])
+
+TOUGH2 = cfg['simulator']['executable']
+BEOPEST = cfg['pest']['executable']
+PESTDIR = cfg['pest']['dir']
+
+"""
+This script helps to launch BeoPEST on the case in current directory (either
+'case.pst' or 'case_svda.pst').  Use:
+
+    python run_beopest.py
+
+Requires:
+    - Python 2.6 or 2.7, NumPy
+    - PEST and TOUGH2 executables, see settings above.
+
+Most of the important settings are listed above.  Instead of letting PEST runs
+TOUGH2 models directly, actual model runs were controlled by the pest_model.py
+script, which uses the goPEST interface and run_ns_pr.py.  (goPEST is a little
+like PLPROC)
+
+In the actual workflow, there are only a couple files that need attention:
+
+    - goPESTpar.list    a list of parameters (goPESTpar.py translates this into
+        actual PEST parameter data)
+    - goPESTobs.list    a list of observations (goPESTobs.py translates this
+        into actual PEST observation data)
+
+Brief workflow:
+
+    1. edit goPESTpar.list and goPESTobs.list
+
+    2. run 'python make_case_pst.py', which processed two lists above and
+    updates case.pst.  This also initialise the parameter values from the values
+    in the real_model_original.dat (.pdat).
+
+    3. inspect/modify/check case.pst
+
+    4. run 'python run_beopest.py'
+
+    5. to obtain the model with optimised parameters run 'python
+    make_real_model.py' or 'python make_real_model_svda.py', now
+    real_model.dat is the updated model, run_ns_pr.py can be used to run    this
+    model.
+
+    6. occasionally run 'make_batch_files.py', eg. after svdaprep, so that model
+    batch files are correct.  This should be run after moving the folder between
+    platform as well.
+
+To update/change a TOUGH2 model, the following files needs to be updated:
+
+    - g_real_model.dat (model's mulgrid/geometry file)
+    - real_model_original.dat (and .pdat)
+    - real_model_original_pr.dat (and .pdat, rocktypes, geners not as important,
+        they are processed from the natural state real_model_ogirinal.dat)
+    - real_model.incon
+    - and of course goPESTpar.list and goPESTobs.list
+    - make sure to run 'python make_case_pst.py'
+
+"""
+
+if NUM_SLAVES is None:
+    NUM_SLAVES = max(cpu_count() - 2, 1)
+if PESTDIR:
+    BEOPEST = os.path.join(PESTDIR,BEOPEST)
+
+def ignore_dirs(folder, files):
+    """ for shutil.copytree, ignores all subdirectories, ie. copy only files
+    """
+    ignore_list = []
+    for f in files:
+        full_path = os.path.join(folder, f)
+        if os.path.isdir(full_path):
+            ignore_list.append(f)
+    return ignore_list
+
+def run_cli(argv=[]):
+    """ generate master and slave commands, copy files into slave directories,
+    and launch them as process, wait until all finished.
+    """
+    if "pest_hp" in BEOPEST.lower():
+        MASTER = BEOPEST
+        if PESTDIR:
+            AGENT = os.path.join(PESTDIR, cfg['pest']['executable_agent'])
+        else:
+            AGENT = cfg['pest']['executable_agent']
+    else:
+        MASTER = BEOPEST
+        AGENT = BEOPEST
+
+    generate_comm_files()
+    print('Running BeoPEST with silent=', SILENT_SLAVES)
+
+    # each command is a tuple (args, other options), see subproces.Popen()
+    master = ([MASTER, PST_NAME, SWITCHES, '/h :%s' % PORT], {} )
+    slaves, slaves_outs = [], []
+    for i in range(NUM_SLAVES):
+        s_dir = 'slave%i' % (i+1)
+        if os.path.exists(s_dir):
+            shutil.rmtree(s_dir)
+        shutil.copytree('.', s_dir, ignore=ignore_dirs)
+        if SILENT_SLAVES:
+            # stdout_file = DEVNULL
+            stdout_file = open(s_dir + '/stdout.txt', 'w')
+            slaves_outs.append(stdout_file)
+        else:
+            stdout_file = None
+        slaves.append((
+            [AGENT, PST_NAME, '/h localhost:%s' % PORT],
+            {
+                'cwd': s_dir,
+                'stdout': stdout_file,
+                'stderr': stdout_file,
+            }))
+
+    # start master, sleep a bit, then start all slaves
+    ps = []
+    ps.append(Popen(master[0], **master[1]))
+    time.sleep(0.1)
+    for s in slaves:
+        ps.append(Popen(s[0], **s[1]))
+    for p in ps:
+        p.wait()
+    for f in slaves_outs:
+        f.close()
+
+def gen_run_management_file(nslave, slaves, wait=0.2, parlam=1, runtime=3600):
+    """ generate run management file required by ppest/jactest.  I am
+    implementing the simplified version here, single template and instruction
+    file only.  TODO extend to be more flexible, which may require more info
+    from pest control file.  Each slaves is a list of tuple (SLAVNAME,SLAVDIR),
+    should have the same length as nslaves.
+    """
+    ifletyp = 0 # for now
+    lines = []
+    lines += [
+        'prf',
+        '%i %i %f %i' % (nslave, ifletyp, wait, parlam),
+    ]
+    for i in range(nslave):
+        slavname, slavdir = slaves[i]
+        if " " in slavname or "'" in slavname:
+            slavname = slavname.replace("'", "''")
+            slavname = "'" + slavname + "'"
+        lines.append(slavname + ' ' + slavdir)
+    lines.append(' '.join([str(runtime)] * nslave))
+    # TODO: support non-zero IFLETYP
+    return '\n'.join(lines)
+
+def run_pslaves(master_command):
+    """ generate master and slave commands, copy files into slave directories,
+    and launch them as process, wait until all finished.  master_command should
+    be a list of strings [cmd, arg1, arg2, ...].
+    """
+    # each command is a tuple (args, other options), see subproces.Popen()
+    master = (master_command, {} )
+    slaves = []
+    for i in range(NUM_SLAVES):
+        s_dir = 'slave%i' % (i+1)
+        if os.path.exists(s_dir):
+            shutil.rmtree(s_dir)
+        shutil.copytree('.', s_dir, ignore=ignore_dirs)
+        slaves.append((
+            [BEOPEST, PST_NAME, '/h localhost:%s' % PORT],
+            {
+                'cwd': s_dir,
+            }))
+
+    # start master, sleep a bit, then start all slaves
+    ps = []
+    ps.append(Popen(master[0], **master[1]))
+    time.sleep(0.1)
+    for s in slaves:
+        ps.append(Popen(s[0], **s[1]))
+    for p in ps:
+        p.wait()
+
+def generate_comm_files():
+    """ generate some files contain information so slaves know where things are.
+    TODO: to improve, there should be better ways, partially to fit with my
+    original NeSI system.
+    """
+    def write_to(filename, line):
+        with open(filename, 'w') as f:
+            f.write(line)
+    write_to('_master_dir', os.getcwd())
+    write_to('_pest_dir', PESTDIR)
+    write_to('_tough2', TOUGH2)
+
+if __name__ == '__main__':
+    run_cli()