[WIP yt 3] Octree ghost zones

[cphyc]

PR Summary

This PR adds support for ghost zones (or a similar interface).

PR Checklist

• Code passes flake8 checker
• New features are documented, with docstrings and narrative docs
• Adds a test for any bugs fixed. Adds tests for new features.
• Face-centred data
• Gradient, vorticity, divergence support

[cphyc]

Using

import yt
import numpy as np

ds = yt.load('output_00080/info_00080.txt')
sp = ds.sphere([.5]*3, (0.5, 'code_length'))
ad = ds.all_data()

def generate_gradient(direction):
    idir = 'xyz'.index(direction)
    def grad(field, data):
        if isinstance(data, yt.fields.field_detector.FieldDetector):
            return data['pressure'] / data['dx']
        offset = 0
        cell_count = 0

        for i, subset in enumerate(data._current_chunk.objs):
            oh = subset.domain.oct_handler
            cell_count += subset.base_selector.count_oct_cells(oh, subset.domain_id)
    
        dest = np.zeros(cell_count, dtype=np.float64)

        for i, subset in enumerate(data._current_chunk.objs):
            oh = subset.domain.oct_handler

            # Extract *all* data in octree
            data_in = {direction: subset[direction],
                       'pressure': subset['pressure']}
            data_out = oh.get_hypercube(subset, data_in)

            sl = slice(1, 3)
            sl1 = [sl]*3 + [slice(None)]
            sl2 = [sl]*3 + [slice(None)]
            sl3 = [sl]*3 + [slice(None)]
            sl1[2-idir] = slice(0, 2)
            sl2[2-idir] = slice(1, 3)
            sl3[2-idir] = slice(2, 4)

            # Compute gradients
            x = data_out[direction]
            p = data_out['pressure']
            dpl = (p[sl3] - p[sl2]) / (x[sl3] - x[sl2])
            dpr = (p[sl2] - p[sl1]) / (x[sl2] - x[sl1])

            maskl = np.isfinite(dpl)
            maskr = np.isfinite(dpr)

            grad = np.where(maskl & maskr, (dpl + dpr) / 2,
                np.where(maskl, dpl, dpr))

            # Extract *all* data in octree
            tmp = grad
            doffset = subset.select(subset.base_selector, tmp, dest, offset)
            offset += doffset

        return data.apply_units(dest, p.units / x.units)
    return grad

def grad_magnitude(field, data):
    values = np.sqrt(data['gas', 'p_grad_x']**2 +
                     data['gas', 'p_grad_y']**2 + 
                     data['gas', 'p_grad_z']**2)
    return values


ds.add_field(('gas', 'p_grad_x'), function=generate_gradient('x'), sampling_type='cell', units='dyne/cm**3')
ds.add_field(('gas', 'p_grad_y'), function=generate_gradient('y'), sampling_type='cell', units='dyne/cm**3')
ds.add_field(('gas', 'p_grad_z'), function=generate_gradient('z'), sampling_type='cell', units='dyne/cm**3')
ds.add_field(('gas', 'p_gradient_magnitude'), function=grad_magnitude, sampling_type='cell', units='dyne/cm**3')

p = yt.SlicePlot(ds, 'x', ['pressure', 'p_gradient_magnitude',
                'p_grad_x', 'p_grad_y', 'p_grad_z'])

we get

Using a covering grid, we get

[ngoldbaum]

I believe you added a file called test.py accidentally.

[cphyc]

@ngoldbaum thanks for the tip! I did indeed (as I'm using it to debug my code). I'll remove it before changing this from WIP to non-WIP.

[matthewturk]

This is pretty rad.

[cphyc]

As of f25e251, here is what we get

Note that the way I did the slice only selects octs in the sliced region, so that the x gradient is not well computed. The y gradients are also fine and consistent with the ones obtained using the covering grid.

---

The script used to generate the files is

import yt
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np

mpl.rcParams['figure.figsize'] = (8, 8)

ds = yt.load('info_00080.txt')
ad = ds.all_data()

def gen_grad(idir):
    def my_grad(field, data):
        if isinstance(data, yt.fields.field_detector.FieldDetector):
            return data['pressure'] / data['dx']
        offset = 0
        cell_count = 0

        for i, subset in enumerate(data._current_chunk.objs):
            oh = subset.domain.oct_handler
            cell_count += subset.base_selector.count_oct_cells(oh, subset.domain_id)

        dest = np.zeros(cell_count, dtype=np.float64)

        for i, subset in enumerate(data._current_chunk.objs):
            oh = subset.domain.oct_handler

            # Extract *all* data in octree
            data_in = {'x': subset['index', 'x'],
                       'dx': subset['index', 'dx'],
                       'pressure': subset['gas', 'pressure']}
            data_out = oh.get_hypercube(subset, data_in)
            cs = slice(1, -1, None)
            slice_base = [cs, cs, cs]
            sl, sc, sr = [slice_base.copy(), slice_base.copy(), slice_base.copy()]
            sl[idir] = slice(0,   -2, None)
            sc[idir] = slice(1,   -1, None)
            sr[idir] = slice(2, None, None)
            sl = tuple(sl)
            sr = tuple(sr)
            sc = tuple(sc)
            p = data_out['pressure']
            dx = data_out['dx']

            ml = np.isfinite(p[sl])
            mr = np.isfinite(p[sr])
            grad_p = np.where(ml & mr, (p[sr] - p[sl]) / 2,
                              np.where(ml, p[sc] - p[sl],
                                       p[sr] - p[sc])) / dx[sc]

            # Extract *all* data in octree
            doffset = subset.select(subset.base_selector, grad_p, dest, offset)
            offset += doffset

        return data.apply_units(dest, p.units / dx.units)

    ds.add_field(('gas', 'grad_p_%s' % 'zyx'[idir]),
                 function=my_grad, units='g/(cm**2*s**2)')
gen_grad(0)
gen_grad(1)
gen_grad(2)

p = yt.SlicePlot(ds, 'x', ['pressure',  'grad_p_z'],
                 # center=center
)
p.set_cmap('all', 'viridis')
f = 'grad_p_z'
p.set_unit(f, 'dyne/cm**3')
p.set_cmap(f, 'bwr')
p.set_zlim(f, -1e-39, 1e-39)
p.set_log(f, True, linthresh=1e-43)

p.save('frames/')

N = 64 * 2**1
grid = ds.smoothed_covering_grid(1, left_edge=[0]*3, dims=[N]*3, num_ghost_zones=1)

p = grid['pressure'].v
pz = grid['pressure_gradient_z'].v

kwa = dict(origin='lower', extent=(grid.left_edge[0], grid.right_edge[0],
                                   grid.left_edge[0], grid.right_edge[0]),
           cmap='bwr', norm=mpl.colors.SymLogNorm(linthresh=1e-43, vmin=-1e-39, vmax=1e-39))

plt.clf()
plt.imshow(pz[N//2].T, **kwa)
plt.colorbar()
plt.savefig('frames/info_00001_Unigrid_pressure_gradient_z.png')

[cphyc]

Interestingly, the following script produces the same plot.

import yt

ds = yt.load('info_00080.txt')
p = p = yt.SlicePlot(ds, 'x', ['pressure', 'pressure_gradient_x'])
p.set_cmap('all', 'viridis')
f = 'pressure_gradient_x'
p.set_unit(f, 'dyne/cm**3')
p.set_cmap(f, 'bwr')
p.set_zlim(f, -1e-39, 1e-39)
p.set_log(f, True, linthresh=1e-43)
p.save()

[EDIT]: the ordering is now correct following 22ef070

[cphyc]

This is superseded by #2166 .