test_particle_filter.py 文件源码

def test_particle_filter_dependency():
    """
    Test dataset add_particle_filter which should automatically add
    the dependency of the filter.
    """

    @particle_filter(filtered_type='all', requires=['particle_type'])
    def stars(pfilter, data):
        filter = data[(pfilter.filtered_type, "particle_type")] == 2
        return filter

    @particle_filter(filtered_type='stars', requires=['creation_time'])
    def young_stars(pfilter, data):
        age = data.ds.current_time - data[pfilter.filtered_type, "creation_time"]
        filter = np.logical_and(age.in_units('Myr') <= 5, age >= 0)
        return filter

    ds = yt.load(iso_galaxy)
    ds.add_particle_filter('young_stars')
    assert 'young_stars' in ds.particle_types
    assert 'stars' in ds.particle_types
    assert ('deposit', 'young_stars_cic') in ds.derived_field_list
    assert ('deposit', 'stars_cic') in ds.derived_field_list