test.py

import yt
import numpy as np

ds = yt.load("enzo_tiny_cosmology/RD0009/RD0009")
ad = ds.all_data()


def _norm(field, data):
    temp = data["density"].copy()
    dens = np.array(data["density"])
    temp.fill(0)
    temp[np.logical_and(dens > 1e-29, dens <= 1e-28)] = 1.0
    temp[np.logical_and(dens > 1e-30, dens <= 1e-29)] = 2.0
    temp[np.logical_and(dens > 1e-31, dens <= 1e-30)] = 3.0
    return temp


ds.add_field("norm", function=_norm, take_log=False, units="g/cm**3")


plot = yt.SlicePlot(ds, "x", "norm")
plot.set_log('norm', False)
plot.set_zlim('norm', 0.0, 3.0)
plot.set_cmap('norm', 'RAINBOW')
plot.save("Test.png")