Compute gradients for octree dataset in yt

octree_gradients.py

def add_deposited_gradient_field(ds, gradient_field, direction, ptype):
    '''Add a gradient field, evaluated at the location of particle

    Parameters
    ----------
    ds : Dataset
    gradient_field : tuple
       The field to compute the gradient of (ftype, fname), e.g. ("gas", "pressure")
    direction : str
       The gradient direction (x, y or z)
    ptype : str
       The name of the particle type (e.g. "DM" or "stars")

    Returns
    -------
    fname : tuple
       The name of the gradient field as a tuple.
    '''
    DIRECTION = 'xyz'.index(direction)
    def _gradient_at_point(field, data):
        ptype, _ = field.name
        # Get the position of the particles
        pos = data[ptype, "particle_position"]
        Npart = pos.shape[0]
        pos_neigh = data.ds.arr((pos, pos))

        if isinstance(data, FieldDetector):
            return np.zeros(Npart)

        dx = data[ptype, f"cell_index_dx"].to('code_length')

        # Compute center of neighboring cells
        pos_neigh[0, :, DIRECTION] -= dx
        pos_neigh[1, :, DIRECTION] += dx
        lvl = data[ptype, 'particle_level'] - data.ds.parameters['levelmin']

        Npart = pos.shape[0]
        tmp = np.zeros(Npart)

        ILEFT = 0
        IRIGHT = 1

        ret = np.zeros((2, Npart))
        remaining = np.ones((2, Npart), dtype=bool)
        Nremaining = [Npart, Npart]

        for i, subset in enumerate(data._current_chunk.objs):
            mesh_data = subset[gradient_field].T.reshape(-1)
            # Get cell index on left side & right side
            for idir in (ILEFT, IRIGHT):
                if Nremaining[idir] == 0:
                    continue
                tmp[:Nremaining[idir]] = subset.mesh_sampling_particle_field(
                    pos_neigh[idir, remaining[idir]].copy(),
                    mesh_data, lvlmax=lvl[remaining[idir]])

                ret[idir, remaining[idir]] = tmp[:Nremaining[idir]]

                remaining[idir, remaining[idir]] = np.isnan(tmp[:Nremaining[idir]])
                Nremaining[idir] = remaining[idir].sum()

        return data.ds.arr((ret[IRIGHT, :] - ret[ILEFT, :]) / dx / 2,
                           mesh_data.units / dx.units)

    units = '%s/%s' % (ds.field_info[gradient_field].units,
                       ds.field_info['index', 'dx'].units)
    fname = 'cell_%s_gradient_%s' % (
        '_'.join(gradient_field),
        direction
        )

    if (ptype, 'cell_index_dx') not in ds.derived_field_list:
        ds.add_mesh_sampling_particle_field(('index', 'dx'), ptype=ptype)

    new_field = (ptype, fname)
    ds.add_field(new_field, function=_gradient_at_point,
                 units=units, sampling_type='particle', force_override=True)

    return new_field


def add_gradient_field(ds, gradient_field, direction):
    '''Add a gradient field

    Parameters
    ----------
    ds : Dataset
    gradient_field : tuple
       The field to compute the gradient of (ftype, fname), e.g. ("gas", "pressure")
    direction : str
       The gradient direction (x, y or z)

    Returns
    -------
    fname : tuple
       The name of the gradient field as a tuple.
    '''
    DIRECTION = 'xyz'.index(direction)
    def _gradient_at_point(field, data):
        # Get the position of the particles
        pos = yt.ustack([data['index', k] for k in 'xyz'], axis=0)
        Npart = pos.shape[0]
        pos_neigh = yt.ustack((pos, pos), axis=-1)

        if isinstance(data, FieldDetector):
            return np.zeros(Npart)

        dx = data['index', 'dx'].to('code_length')

        # Compute center of neighboring cells
        pos_neigh[0, :, DIRECTION] -= dx
        pos_neigh[1, :, DIRECTION] += dx
        lvl = data['index', 'grid_level']

        Npart = pos.shape[0]
        tmp = np.zeros(Npart)

        ILEFT = 0
        IRIGHT = 1

        ret = np.zeros((2, Npart))
        remaining = np.ones((2, Npart), dtype=bool)
        Nremaining = [Npart, Npart]

        for i, subset in enumerate(data._current_chunk.objs):
            mesh_data = subset[gradient_field].T.reshape(-1)
            # Get cell index on left side & right side
            for idir in (ILEFT, IRIGHT):
                if Nremaining[idir] == 0:
                    continue
                tmp[:Nremaining[idir]] = subset.mesh_sampling_particle_field(
                    pos_neigh[idir, remaining[idir]].copy(),
                    mesh_data, lvlmax=lvl[remaining[idir]])

                ret[idir, remaining[idir]] = tmp[:Nremaining[idir]]

                remaining[idir, remaining[idir]] = np.isnan(tmp[:Nremaining[idir]])
                Nremaining[idir] = remaining[idir].sum()

        return data.ds.arr((ret[IRIGHT, :] - ret[ILEFT, :]) / dx / 2,
                           mesh_data.units / dx.units)

    units = '%s/%s' % (ds.field_info[gradient_field].units,
                       ds.field_info['index', 'dx'].units)

    ptype, fname = gradient_field
    new_field = (ptype, f'{fname}_gradient_{direction}')

    ds.add_field(new_field, function=_gradient_at_point,
                 units=units, sampling_type='cell', force_override=True)

    return new_field