Combining rotation and SPH calculations (v_disp, rho, smooth) is broken

[Martin-Rey]

Hi,

I am coming across an unexpected behaviour when combining snapshot rotations (e.g. through faceon) and SPH calculations (e.g. through derivations of the local velocity dispersion v_disp). A minimum, breaking example with Traceback is available below.

The reason is somewhat obvious, even if it took me a while to understand it. Once a rotation has been applied, particle positions x, y, z are no longer within the cube delimited by the box size, i.e. assert ((snap.d['x'].max() - snap.d['x'].min()) < snap.properties['boxsize'].in_units("kpc")) is False. This then prevents us to derive SPH calculations across the ancestor snapshot with the failure:

ValueError: The particles span a region larger than the specified boxsize

I didn't expect this to be an issue when working with the halo particles only (i.e. the sub snap). But v_disp seems to always be computed at the ancestor snapshot level, even when asking for halo.d['v_disp'].

Furthermore, deriving v_disp across the entire box first, and then doing the halo centering and rotating also doesn't work -- for some reason, the halo sub snap particles don't inherit the v_disp derived array, they attempt to rederive it on their own subsnap, which then asks for the ancestor calculation again, which then fails because of the rotated particles.

This implies that the KDtree and SPH calculations can't handle rotated snapshots because the check for the boxsize here https://github.com/pynbody/pynbody/blob/e640ba32f8beb6e47ccb0c1db5527a5502d05224/pynbody/sph/smooth.cpp#L41-L46 is based on particles positions rather than distances. (I have checked that rotating a snapshot and asking for rho also breaks).

I am unclear on what to do or what to fix to avoid this behaviour. Loosening the boxsize check feels a bit dangerous, so maybe the key is in assuring that v_disp gets inherited correctly onto sub-snaps, or is calculated only from the local particles of the subsnap rather than the full ancestor as is currently done for rho.

Any insights or help would be much appreciated :)

Martin

Breaking example on testdata

    f = pynbody.load("./testdata/g15784.lr.01024.gz")
    f.physical_units()
    h = f.halos()[1]
    transform_chain = pynbody.analysis.angmom.faceon(h.g, disk_size="1 kpc", cen_size="1 kpc")
    h.d['v_disp']      # Also breaks with gas and stars, or when working with the ancestor f.d, f.st, f.g, or when asking for f.d['rho']

failing with

 >       h.d['v_disp']

transformation_test.py:139: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:258: in __getitem__
    return self._get_array_with_lazy_actions(i)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:358: in _get_array_with_lazy_actions
    self.__derive_if_required(name)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:373: in __derive_if_required
    self._derive_array(name)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:1948: in _derive_array
    self.base._derive_array(array_name, self._unifamily)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:1736: in _derive_array
    self.base._derive_array(array_name, fam)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:1433: in _derive_array
    result = fn(self[fam])
../../../.python3_base/lib/python3.9/site-packages/pynbody/derived.py:160: in v_disp
    return _v_sph_operation(self, "disp")
../../../.python3_base/lib/python3.9/site-packages/pynbody/derived.py:137: in _v_sph_operation
    self.kdtree.set_array_ref('rho', self['rho'])
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:258: in __getitem__
    return self._get_array_with_lazy_actions(i)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:358: in _get_array_with_lazy_actions
    self.__derive_if_required(name)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:373: in __derive_if_required
    self._derive_array(name)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:1948: in _derive_array
    self.base._derive_array(array_name, self._unifamily)
../../../.python3_base/lib/python3.9/site-packages/pynbody/snapshot/__init__.py:1433: in _derive_array
    result = fn(self[fam])
../../../.python3_base/lib/python3.9/site-packages/pynbody/sph/__init__.py:172: in rho
    self.kdtree.set_array_ref('smooth',_get_smooth_array_ensuring_compatibility(self))
../../../.python3_base/lib/python3.9/site-packages/pynbody/sph/__init__.py:155: in _get_smooth_array_ensuring_compatibility
    self['smooth'] = smooth_ar = smooth(self)
../../../.python3_base/lib/python3.9/site-packages/pynbody/sph/__init__.py:139: in smooth
    self.kdtree.populate('hsm', config['sph']['smooth-particles'], config['sph']['Kernel'])
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <pynbody.sph.kdtree.KDTree object at 0x13abe0dc0>, mode = 'hsm', nn = 32
kernel = 'CubicSpline'

    def populate(self, mode, nn, kernel = 'CubicSpline'):
        """Create the KDTree and perform the operation specified by `mode`.

        Parameters
        --------
        mode : str (see `kdtree.smooth_operation_to_id`)
            Specify operation to perform (compute smoothing lengths, density, SPH mean, or SPH dispersion).
        nn : int
            Number of neighbours to be considered when smoothing.
        kernel : str
            Keyword to specify the smoothing kernel. Options: 'CubicSpline', 'WendlandC2'
        """
        from . import _thread_map

        n_proc = config["number_of_threads"]

        if kdmain.has_threading() is False and n_proc > 1:
            n_proc = 1
            warnings.warn(
                "Pynbody is configured to use threading for the KDTree, but pthread support was not available during compilation. Reverting to single thread.",
                RuntimeWarning,
            )

        if nn is None:
            nn = 64

>       smx = kdmain.nn_start(self.kdtree, int(nn), n_proc, self.boxsize)
E       ValueError: The particles span a region larger than the specified boxsize

[apontzen]

I agree it's an annoying edge-case. At a practical level, one can simply wrap the box after applying the transformation (and then wrap it again when undoing the transformation, if desired). [EDIT: that's probably not a good idea, see alternative solution below] However it's unlikely a new user would guess this. Perhaps a more helpful error message might be the simple fix? (i.e. suggest within the error message that the user call .wrap() and tries the operation again)

[5p4k]

I have encounered this while helping out some student; in this case, we are calling sideon (with move_all=False to save memory, I believe that is the culprit) and then plotting. It results in the same error.

I have a superfiicial knowledge of pynbody, if there is a workaround possible or a hack, I would be interested to know it so we can move past this roadblock (until a proper fix is posted). @apontzen if it's easy, do you mind sharing a snippet?

On a side node, the kdtree might need to have lower and upper bounds to all coordinates to properly allocate its data, so the exception could be appropriate there.

[apontzen]

I think the best workaround for your case, if you specifically want to use move_all=False, would be to artificially trigger the smoothing operation before you perform the edgeon, i.e. something like:

f.gas['rho']
f.gas['v_disp']
# also access whatever else might be needed...

with pynbody.analysis.angmom.edgeon(h[1], move_all=False):
   ...

I am not sure that move_all=False will save any memory, btw... it might be very slightly quicker, but I doubt it makes much difference by any metric.

[5p4k]

Thank you very much @apontzen for the quick answer, I will try and post back the result!

---

As far as the move_all goes: with move_all=True (default), the script crashed on a small machine with this error: https://github.com/pynbody/pynbody/blob/6b0cb61d89b17451f3bf65866e8891f2aefc4453/pynbody/array.py#L1035

I tracked how much it had tried to allocate, and it had tried to allocate 20GB of data (which failed). With move_all=False, not only it went through, but the total RAM usage for the script did not go above 8GB. I'm not sure how exactly that is the case, but I presume that not all operations of Transformation and its subclasses are occurring in-place, and some temporary has to be allocated to hold the result; the temporary array size is then much smaller when move_all=False.

[5p4k]

After a bit of digging, I managed to get this to work thanks to @apontzen's workaround.

The smoothing is triggered here: https://github.com/pynbody/pynbody/blob/6b0cb61d89b17451f3bf65866e8891f2aefc4453/pynbody/sph/__init__.py#L811

The snipped I was working on is roughly as follows:

data = pynbody.load('...')
stellar_data = data.stars
cluster_data = stellar_data[pynbody.filt.Sphere(radius, centre)]

# Here we need to trigger the smoothing *before* sideon
trigger_cache_generation(cluster_data)

with angmom.sideon(cluster_data, move_all=False):
    pynbody.plot.image(cluster_data, width=2.1 * radius, filename='sideon.png')

and by doing the same caching operation before sideon is called, the error disappears:

def trigger_cache_generation(sim):
    for attr_name in ('x', 'y', 'z', 'pos', 'smooth', 'rho', 'mass'):
        _ = sim[attr_name]

I believe that this issue is in fact the same problem I encountered and this workaround should apply there too, I'll cross-post this solution there in case someone needs it. Thank you all for your input!