Source

bluewaters.support_files / user_script.py

The default branch has multiple heads

Full commit
from yt.config import ytcfg; ytcfg["yt","loglevel"]="50"
from yt.mods import * # Will need to be replaced if we use non-GNU for Enzo
from yt.frontends.enzo.api import EnzoStaticOutputInMemory
from yt.utilities.amr_utils import QuadTree
import h5py
import time
import gc
from contextlib import contextmanager

from mpi4py import MPI
mpi_proc = MPI.COMM_WORLD.rank
mpi_size = MPI.COMM_WORLD.size
root = (mpi_proc == 0)

# (Adaptive) Projections "Field (weight field)"
#   * Density (Density)
#   * Density (None)
#   * Temperature (Density)
#   * HII_Fraction (Density)
#   * H2I_Fraction (Density)
#   * Metal_Density (Density)
#   * HII_Density (None)
#   * HeI_Density (None)
#   * Metal_Density (None)
# 1D Profiles "f1, f2 (weight)"
#   * Density, Temperature (CellMassMsun)
#   * Density, CellMassMsun (None)
#   * Density, Metal_Density (CellMassMsun)
#   * Density, HII_Fraction (CellMassMsun)
# 2D Profiles "f1, f2, f3 (weight)"
#   * Density, Temperature, CellMassMsun (None)
#   * Density, Temperature, MetalMassMsun (None)
#   * Density, Temperature, Metallicity, CellMassMsun
#   * Density, MetalDensity, CellMassMsun (None)
#   * Density, Metallicity, CellMassMsun

projs = (
    ( ["Density", "Temperature", "HII_Fraction",
       "H2I_Fraction", "Metal_Density"],
      "Density"),
    ( ["Density", "HII_Density", "HeI_Density", "Metal_Density", "Ones"],
      None),
      )

prof_1d = ( ("Density", "Temperature", "CellMassMsun"),
            ("Density", "CellMassMsun", None),
            ("Density", "Metal_Density", "CellMassMsun"),
            ("Density", "HII_Fraction", "CellMassMsun")
          )

prof_2d = ( ("Density", "Temperature", "CellMassMsun", None),
            ("Density", "Temperature", "Metal_MassMsun", None),
            ("Density", "Temperature", "Metallicity", "CellMassMsun"),
            ("Density", "Metal_Density", "CellMassMsun", None),
            ("Density", "Metallicity", "CellMassMsun", None),
          )

ex = ["Density", "Temperature", "Metal_Density", "Metallicity"]

def write_profiles(f, p):
    if len(p.field_data["UsedBins"].shape) == 1:
        name = "/1D_%s" % (p.bin_field)
    elif len(p.field_data["UsedBins"].shape) == 2:
        name = "/2D_%s_%s" % (p.x_bin_field, p.y_bin_field)
    g = f.create_group(name)
    for ds_name in p.field_data:
        g.create_dataset(ds_name, data=p[ds_name])

def new_profile(dd, fields, extrema):
    if len(fields) == 1:
        f = fields[0]
        p = BinnedProfile1D(dd, 128, f, extrema[f][0], extrema[f][1],
                            lazy_reader=True)
    elif len(fields) == 2:
        f1 = fields[0]
        f2 = fields[1]
        p = BinnedProfile2D(dd, 128, f1, extrema[f1][0], extrema[f1][1], True,
                                128, f2, extrema[f2][0], extrema[f2][1], True,
                                lazy_reader=True)
    return p

@contextmanager
def time_region(header):
    if root:
        print "YTINLINE: Initiating %s (%0.3e mb)" % (
                    header, get_memory_usage())
    t1 = time.time()
    yield
    t2 = time.time()
    if root:
        print "YTINLINE: Finalizing %s (%0.3e mb, %0.3es)" % (
                    header, get_memory_usage(), t2-t1)

def main():
    t1 = time.time()
    with time_region("Initializing"):
        pf = EnzoStaticOutputInMemory()
        pf.h
    if root and not os.path.exists("%s" % pf):
        os.makedirs("%s" % pf)
    MPI.COMM_WORLD.Barrier()
    with time_region("Projections (weighted)"):
        non_reduced_proj(pf, projs[0][0], projs[0][1], "weighted")
    with time_region("Projections (unweighted)"):
        non_reduced_proj(pf, projs[1][0], projs[1][1], "unweighted")
    # Now we get our profiles
    # First, figure out our extrema ...
    dd = pf.h.all_data()
    with time_region("Extrema"):
        v = dd.quantities["Extrema"](ex)
    extrema = dict(zip(ex, v))

    with time_region("1D Profiles part 1"):
        p1 = new_profile(dd, ["Density"], extrema)
        p1.add_fields(["Temperature", "Metal_Density", "HII_Fraction"],
                      weight = "CellMassMsun")
        p1.add_fields(["CellMassMsun"], weight = None)
    # Save it out
    if root:
        f = h5py.File("%s/profiles.h5" % pf, "w")
        write_profiles(f, p1)

    with time_region("2D Profiles part 1"):
        p2 = new_profile(dd, ["Density", "Temperature"], extrema)
        p2.add_fields(["CellMassMsun", "Metal_MassMsun"], weight=None)
        p2.add_fields(["Metallicity"], weight="CellMassMsun")
    if root: write_profiles(f, p2)

    with time_region("2D Profiles part 2"):
        p3 = new_profile(dd, ["Density", "Metal_Density"], extrema)
        p3.add_fields(["CellMassMsun"], weight=None)
    if root: write_profiles(f, p3)
    with time_region("2D Profiles part 3"):
        p4 = new_profile(dd, ["Density", "Metallicity"], extrema)
        p4.add_fields(["CellMassMsun"], weight=None)
    if root:
        print "YTINLINE: Writing profiles to disk"
        write_profiles(f, p4)
        for exn, (mi, ma) in extrema.items():
            f["/"].attrs["%s_min" % exn] = mi
            f["/"].attrs["%s_max" % exn] = ma
        f.close()
    if root: print "YTINLINE: Concluded. (%0.3e)" % (time.time() - t1)
    for g in pf.h.grids: g.clear_data()
    del pf.h.io, pf, p1, p2, p3, p4
    nobj = gc.collect()
    if root:
        print "YTINLINE: Collected %s objects (%0.3e mb)" % (
            nobj, get_memory_usage())

def add_grid_to_tree(tree, g, fields, weight, dx, axis):
    if weight is None:
        weight_data = g.child_mask.astype("float64") # copy
    else:
        weight_data = g[weight] * g.child_mask
    field_proj = [(g[field] * weight_data).sum(axis=axis) * dx
                  for field in fields]
    weight_proj = weight_data.sum(axis=axis) * dx
    xind, yind = [a.ravel() for a in na.indices(weight_proj.shape)]
    start_index = g.get_global_startindex()
    xind += start_index[x_dict[axis]].astype("int64")
    yind += start_index[y_dict[axis]].astype("int64")
    field_proj = na.array([a.ravel() for a in field_proj], order='F')
    weight_proj = weight_proj.ravel()
    tree.add_array_to_tree(g.Level, xind, yind, field_proj, weight_proj)

def non_reduced_proj(pf, fields, weight, fn_prefix):
    my_proc = (pf.h.grid_procs == mpi_proc)
    fields = ensure_list(fields)
    fn = "%s/%s_%05i.h5" % (pf, fn_prefix, mpi_proc)
    f = h5py.File(fn, "w")
    if root: print "OPENED ", fn
    for ax in range(3):
        xd = pf.domain_dimensions[x_dict[ax]]
        yd = pf.domain_dimensions[y_dict[ax]]
        md = pf.domain_dimensions[ax]
        base_dx = pf.domain_width[ax] / md
        dims = na.array([xd,yd], dtype='int64')
        tree = QuadTree(dims, len(fields), style="integrate")
        for level in range(pf.h.max_level + 1):
            my_level = (pf.h.grid_levels == level)
            gs = pf.h.grids[(my_level & my_proc).ravel()]
            if gs.size == 0: continue
            dl = base_dx/(pf.refine_by**level)
            for g in gs:
                add_grid_to_tree(tree, g, fields, weight, dl, ax)
        buf = tree.tobuffer()
        f.create_group("axis_%02i" % (ax))
        f.create_dataset("axis_%02i/refined" % (ax), data=buf[0])
        f.create_dataset("axis_%02i/values" % (ax), data=buf[1])
        f.create_dataset("axis_%02i/weight" % (ax), data=buf[2])
        f.flush()
    f.create_group("/attributes")
    f["/attributes"].attrs["fields"] = fields
    f["/attributes"].attrs["weights"] = str(weight)
    f["/attributes"].attrs["pf_name"] = str(pf)
    f["/attributes"].attrs["redshift"] = pf.current_redshift
    f["/attributes"].attrs["time_seconds"] = pf.current_time * pf["Time"]
    f.flush()
    f.close()
    del f