Source code for yt.visualization.volume_rendering.transfer_functions

import numpy as np
from more_itertools import always_iterable

from yt.funcs import mylog
from yt.utilities.physical_constants import clight, hcgs, kboltz


[docs] class TransferFunction: r"""A transfer function governs the transmission of emission and absorption through a volume. Transfer functions are defined by boundaries, bins, and the value that governs transmission through that bin. This is scaled between 0 and 1. When integrating through a volume the value through a given cell is defined by the value calculated in the transfer function. Parameters ---------- x_bounds : tuple of floats The min and max for the transfer function. Values below or above these values are discarded. nbins : int How many bins to calculate; in between, linear interpolation is used, so low values are typically fine. Notes ----- Typically, raw transfer functions are not generated unless particular and specific control over the integration is desired. Usually either color transfer functions, where the color values are calculated from color tables, or multivariate transfer functions are used. """ def __init__(self, x_bounds, nbins=256): self.pass_through = 0 self.nbins = nbins # Strip units off of x_bounds, if any x_bounds = [np.float64(xb) for xb in x_bounds] self.x_bounds = x_bounds self.x = np.linspace(x_bounds[0], x_bounds[1], nbins, dtype="float64") self.y = np.zeros(nbins, dtype="float64") self.grad_field = -1 self.light_source_v = self.light_source_c = np.zeros(3, "float64") self.features = []
[docs] def add_gaussian(self, location, width, height): r"""Add a Gaussian distribution to the transfer function. Typically, when rendering isocontours, a Gaussian distribution is the easiest way to draw out features. The spread provides a softness. The values are calculated as :math:`f(x) = h \exp{-(x-x_0)^2 / w}`. Parameters ---------- location : float The centroid of the Gaussian (:math:`x_0` in the above equation.) width : float The relative width (:math:`w` in the above equation.) height : float The peak height (:math:`h` in the above equation.) Note that while values greater 1.0 will be accepted, the values of the transmission function are clipped at 1.0. Examples -------- >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-9.0, 0.01, 1.0) """ vals = height * np.exp(-((self.x - location) ** 2.0) / width) self.y = np.clip(np.maximum(vals, self.y), 0.0, np.inf) self.features.append( ( "gaussian", f"location(x):{location:3.2g}", f"width(x):{width:3.2g}", f"height(y):{height:3.2g}", ) )
[docs] def add_line(self, start, stop): r"""Add a line between two points to the transmission function. This will accept a starting point in (x,y) and an ending point in (x,y) and set the values of the transmission function between those x-values to be along the line connecting the y values. Parameters ---------- start : tuple of floats (x0, y0), the starting point. x0 is between the bounds of the transfer function and y0 must be between 0.0 and 1.0. stop : tuple of floats (x1, y1), the ending point. x1 is between the bounds of the transfer function and y1 must be between 0.0 and 1.0. Examples -------- This will set the transfer function to be linear from 0.0 to 1.0, across the bounds of the function. >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_line((-10.0, 0.0), (-5.0, 1.0)) """ x0, y0 = start x1, y1 = stop slope = (y1 - y0) / (x1 - x0) # We create a whole new set of values and then backout the ones that do # not satisfy our bounding box arguments vals = slope * (self.x - x0) + y0 vals[~((self.x >= x0) & (self.x <= x1))] = 0.0 self.y = np.clip(np.maximum(vals, self.y), 0.0, np.inf) self.features.append( ( "line", f"start(x,y):({start[0]:3.2g}, {start[1]:3.2g})", f"stop(x,y):({stop[0]:3.2g}, {stop[1]:3.2g})", ) )
[docs] def add_step(self, start, stop, value): r"""Adds a step function to the transfer function. This accepts a `start` and a `stop`, and then in between those points the transfer function is set to the maximum of the transfer function and the `value`. Parameters ---------- start : float This is the beginning of the step function; must be within domain of the transfer function. stop : float This is the ending of the step function; must be within domain of the transfer function. value : float The value the transfer function will be set to between `start` and `stop`. Note that the transfer function will *actually* be set to max(y, value) where y is the existing value of the transfer function. Examples -------- Note that in this example, we have added a step function, but the Gaussian that already exists will "win" where it exceeds 0.5. >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-7.0, 0.01, 1.0) >>> tf.add_step(-8.0, -6.0, 0.5) """ vals = np.zeros(self.x.shape, "float64") vals[(self.x >= start) & (self.x <= stop)] = value self.y = np.clip(np.maximum(vals, self.y), 0.0, np.inf) self.features.append( ( "step", f"start(x):{start:3.2g}", f"stop(x):{stop:3.2g}", f"value(y):{value:3.2g}", ) )
[docs] def add_filtered_planck(self, wavelength, trans): from yt._maintenance.numpy2_compat import trapezoid vals = np.zeros(self.x.shape, "float64") nu = clight / (wavelength * 1e-8) nu = nu[::-1] for i, logT in enumerate(self.x): T = 10**logT # Black body at this nu, T Bnu = ((2.0 * hcgs * nu**3) / clight**2.0) / ( np.exp(hcgs * nu / (kboltz * T)) - 1.0 ) # transmission f = Bnu * trans[::-1] # integrate transmission over nu vals[i] = trapezoid(f, nu) # normalize by total transmission over filter self.y = vals / trans.sum()
# self.y = np.clip(np.maximum(vals, self.y), 0.0, 1.0)
[docs] def plot(self, filename): r"""Save an image file of the transfer function. This function loads up matplotlib, plots the transfer function and saves. Parameters ---------- filename : string The file to save out the plot as. Examples -------- >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-9.0, 0.01, 1.0) >>> tf.plot("sample.png") """ import matplotlib import matplotlib.pyplot as plt matplotlib.use("Agg") plt.clf() plt.plot(self.x, self.y, "xk-") plt.xlim(*self.x_bounds) plt.ylim(0.0, 1.0) plt.savefig(filename)
[docs] def show(self): r"""Display an image of the transfer function This function loads up matplotlib and displays the current transfer function. Parameters ---------- Examples -------- >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-9.0, 0.01, 1.0) >>> tf.show() """ import matplotlib.pyplot as plt plt.clf() plt.plot(self.x, self.y, "xk-") plt.xlim(*self.x_bounds) plt.ylim(0.0, 1.0) plt.draw()
[docs] def clear(self): self.y[:] = 0.0 self.features = []
def __repr__(self): disp = ( "<Transfer Function Object>: " f"x_bounds:({self.x_bounds[0]:3.2g}, {self.x_bounds[1]:3.2g}) " f"nbins:{self.nbins:3.2g} features:{self.features}" ) return disp
[docs] class MultiVariateTransferFunction: r"""This object constructs a set of field tables that allow for multiple field variables to control the integration through a volume. The integration through a volume typically only utilizes a single field variable (for instance, Density) to set up and control the values returned at the end of the integration. For things like isocontours, this is fine. However, more complicated schema are possible by using this object. For instance, density-weighted emission that produces colors based on the temperature of the fluid. Parameters ---------- grey_opacity : bool Should opacity be calculated on a channel-by-channel basis, or overall? Useful for opaque renderings. Default: False """ def __init__(self, grey_opacity=False): self.n_field_tables = 0 self.tables = [] # Tables are interpolation tables self.field_ids = [0] * 6 # This correlates fields with tables self.weight_field_ids = [-1] * 6 # This correlates self.field_table_ids = [0] * 6 self.weight_table_ids = [-1] * 6 self.grad_field = -1 self.light_source_v = self.light_source_c = np.zeros(3, "float64") self.grey_opacity = grey_opacity
[docs] def add_field_table(self, table, field_id, weight_field_id=-1, weight_table_id=-1): r"""This accepts a table describing integration. A "field table" is a tabulated set of values that govern the integration through a given field. These are defined not only by the transmission coefficient, interpolated from the table itself, but the `field_id` that describes which of several fields the integration coefficient is to be calculated from. Parameters ---------- table : `TransferFunction` The integration table to be added to the set of tables used during the integration. field_id : int Each volume has an associated set of fields. This identifies which of those fields will be used to calculate the integration coefficient from this table. weight_field_id : int, optional If specified, the value of the field this identifies will be multiplied against the integration coefficient. weight_table_id : int, optional If specified, the value from the *table* this identifies will be multiplied against the integration coefficient. Notes ----- This can be rather complicated. It's recommended that if you are interested in manipulating this in detail that you examine the source code, specifically the function FIT_get_value in yt/_amr_utils/VolumeIntegrator.pyx. Examples -------- This example shows how to link a new transfer function against field 0. Note that this by itself does not link a *channel* for integration against a field. This is because the weighting system does not mandate that all tables contribute to a channel, only that they contribute a value which may be used by other field tables. >>> mv = MultiVariateTransferFunction() >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-7.0, 0.01, 1.0) >>> mv.add_field_table(tf, 0) """ self.tables.append(table) self.field_ids[self.n_field_tables] = field_id self.weight_field_ids[self.n_field_tables] = weight_field_id self.weight_table_ids[self.n_field_tables] = weight_table_id self.n_field_tables += 1
[docs] class ColorTransferFunction(MultiVariateTransferFunction): r"""A complete set of transfer functions for standard color-mapping. This is the best and easiest way to set up volume rendering. It creates field tables for all three colors, their alphas, and has support for sampling color maps and adding independent color values at all locations. It will correctly set up the `MultiVariateTransferFunction`. Parameters ---------- x_bounds : tuple of floats The min and max for the transfer function. Values below or above these values are discarded. nbins : int How many bins to calculate; in between, linear interpolation is used, so low values are typically fine. grey_opacity : bool Should opacity be calculated on a channel-by-channel basis, or overall? Useful for opaque renderings. """ def __init__(self, x_bounds, nbins=256, grey_opacity=False): MultiVariateTransferFunction.__init__(self) # Strip units off of x_bounds, if any x_bounds = [np.float64(xb) for xb in x_bounds] self.x_bounds = x_bounds self.nbins = nbins # This is all compatibility and convenience. self.red = TransferFunction(x_bounds, nbins) self.green = TransferFunction(x_bounds, nbins) self.blue = TransferFunction(x_bounds, nbins) self.alpha = TransferFunction(x_bounds, nbins) self.funcs = (self.red, self.green, self.blue, self.alpha) self.grey_opacity = grey_opacity self.features = [] # Now we do the multivariate stuff # We assign to Density, but do not weight for i, tf in enumerate(self.funcs[:3]): self.add_field_table(tf, 0, weight_table_id=3) self.link_channels(i, i) self.add_field_table(self.funcs[3], 0) self.link_channels(3, 3) # We don't have a fifth table, so the value will *always* be zero. # self.link_channels(4, [3,4,5])
[docs] def add_gaussian(self, location, width, height): r"""Add a Gaussian distribution to the transfer function. Typically, when rendering isocontours, a Gaussian distribution is the easiest way to draw out features. The spread provides a softness. The values are calculated as :math:`f(x) = h \exp{-(x-x_0)^2 / w}`. Parameters ---------- location : float The centroid of the Gaussian (:math:`x_0` in the above equation.) width : float The relative width (:math:`w` in the above equation.) height : list of 4 float The peak height (:math:`h` in the above equation.) Note that while values greater 1.0 will be accepted, the values of the transmission function are clipped at 1.0. This must be a list, and it is in the order of (red, green, blue, alpha). Examples -------- This adds a red spike. >>> tf = ColorTransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-9.0, 0.01, [1.0, 0.0, 0.0, 1.0]) """ for tf, v in zip(self.funcs, height): tf.add_gaussian(location, width, v) self.features.append( ( "gaussian", f"location(x):{location:3.2g}", f"width(x):{width:3.2g}", f"height(y):({height[0]:3.2g}, {height[1]:3.2g}, {height[2]:3.2g}, {height[3]:3.2g})", ) )
[docs] def add_step(self, start, stop, value): r"""Adds a step function to the transfer function. This accepts a `start` and a `stop`, and then in between those points the transfer function is set to the maximum of the transfer function and the `value`. Parameters ---------- start : float This is the beginning of the step function; must be within domain of the transfer function. stop : float This is the ending of the step function; must be within domain of the transfer function. value : list of 4 floats The value the transfer function will be set to between `start` and `stop`. Note that the transfer function will *actually* be set to max(y, value) where y is the existing value of the transfer function. This must be a list, and it is in the order of (red, green, blue, alpha). Examples -------- This adds a step function that will produce a white value at > -6.0. >>> tf = ColorTransferFunction((-10.0, -5.0)) >>> tf.add_step(-6.0, -5.0, [1.0, 1.0, 1.0, 1.0]) """ for tf, v in zip(self.funcs, value): tf.add_step(start, stop, v) self.features.append( ( "step", f"start(x):{start:3.2g}", f"stop(x):{stop:3.2g}", f"value(y):({value[0]:3.2g}, {value[1]:3.2g}, {value[2]:3.2g}, {value[3]:3.2g})", ) )
[docs] def plot(self, filename): r"""Save an image file of the transfer function. This function loads up matplotlib, plots all of the constituent transfer functions and saves. Parameters ---------- filename : string The file to save out the plot as. Examples -------- >>> tf = ColorTransferFunction((-10.0, -5.0)) >>> tf.add_layers(8) >>> tf.plot("sample.png") """ from matplotlib import pyplot from matplotlib.ticker import FuncFormatter pyplot.clf() ax = pyplot.axes() i_data = np.zeros((self.alpha.x.size, self.funcs[0].y.size, 3)) i_data[:, :, 0] = np.outer(np.ones(self.alpha.x.size), self.funcs[0].y) i_data[:, :, 1] = np.outer(np.ones(self.alpha.x.size), self.funcs[1].y) i_data[:, :, 2] = np.outer(np.ones(self.alpha.x.size), self.funcs[2].y) ax.imshow(i_data, origin="lower") ax.fill_between( np.arange(self.alpha.y.size), self.alpha.x.size * self.alpha.y, y2=self.alpha.x.size, color="white", ) ax.set_xlim(0, self.alpha.x.size) xticks = ( np.arange(np.ceil(self.alpha.x[0]), np.floor(self.alpha.x[-1]) + 1, 1) - self.alpha.x[0] ) xticks *= (self.alpha.x.size - 1) / (self.alpha.x[-1] - self.alpha.x[0]) ax.xaxis.set_ticks(xticks) def x_format(x, pos): return "%.1f" % ( x * (self.alpha.x[-1] - self.alpha.x[0]) / (self.alpha.x.size - 1) + self.alpha.x[0] ) ax.xaxis.set_major_formatter(FuncFormatter(x_format)) yticks = np.linspace(0, 1, 5) * self.alpha.y.size ax.yaxis.set_ticks(yticks) def y_format(y, pos): return y / self.alpha.y.size ax.yaxis.set_major_formatter(FuncFormatter(y_format)) ax.set_ylabel("Transmission") ax.set_xlabel("Value") pyplot.savefig(filename)
[docs] def show(self, ax=None): r"""Display an image of the transfer function This function loads up matplotlib and displays the current transfer function. Parameters ---------- Examples -------- >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-9.0, 0.01, 1.0) >>> tf.show() """ from matplotlib import pyplot from matplotlib.ticker import FuncFormatter pyplot.clf() ax = pyplot.axes() i_data = np.zeros((self.alpha.x.size, self.funcs[0].y.size, 3)) i_data[:, :, 0] = np.outer(np.ones(self.alpha.x.size), self.funcs[0].y) i_data[:, :, 1] = np.outer(np.ones(self.alpha.x.size), self.funcs[1].y) i_data[:, :, 2] = np.outer(np.ones(self.alpha.x.size), self.funcs[2].y) ax.imshow(i_data, origin="lower") ax.fill_between( np.arange(self.alpha.y.size), self.alpha.x.size * self.alpha.y, y2=self.alpha.x.size, color="white", ) ax.set_xlim(0, self.alpha.x.size) xticks = ( np.arange(np.ceil(self.alpha.x[0]), np.floor(self.alpha.x[-1]) + 1, 1) - self.alpha.x[0] ) xticks *= (self.alpha.x.size - 1) / (self.alpha.x[-1] - self.alpha.x[0]) if len(xticks) > 5: xticks = xticks[:: len(xticks) // 5] ax.xaxis.set_ticks(xticks) def x_format(x, pos): return "%.1f" % ( x * (self.alpha.x[-1] - self.alpha.x[0]) / (self.alpha.x.size - 1) + self.alpha.x[0] ) ax.xaxis.set_major_formatter(FuncFormatter(x_format)) yticks = np.linspace(0, 1, 5) * self.alpha.y.size ax.yaxis.set_ticks(yticks) def y_format(y, pos): s = f"{y:0.2f}" return s ax.yaxis.set_major_formatter(FuncFormatter(y_format)) ax.set_ylabel("Opacity") ax.set_xlabel("Value")
[docs] def vert_cbar( self, resolution, log_scale, ax, label=None, label_fmt=None, *, label_fontsize=10, size=10, ): r"""Display an image of the transfer function This function loads up matplotlib and displays the current transfer function. Parameters ---------- Examples -------- >>> tf = TransferFunction((-10.0, -5.0)) >>> tf.add_gaussian(-9.0, 0.01, 1.0) >>> tf.show() """ from matplotlib.ticker import FuncFormatter if label is None: label = "" alpha = self.alpha.y max_alpha = alpha.max() i_data = np.zeros((self.alpha.x.size, self.funcs[0].y.size, 3)) i_data[:, :, 0] = np.outer(self.funcs[0].y, np.ones(self.alpha.x.size)) i_data[:, :, 1] = np.outer(self.funcs[1].y, np.ones(self.alpha.x.size)) i_data[:, :, 2] = np.outer(self.funcs[2].y, np.ones(self.alpha.x.size)) ax.imshow(i_data, origin="lower", aspect="auto") ax.plot(alpha, np.arange(self.alpha.y.size), "w") # Set TF limits based on what is visible visible = np.argwhere(self.alpha.y > 1.0e-3 * self.alpha.y.max()) # Display colobar values xticks = ( np.arange(np.ceil(self.alpha.x[0]), np.floor(self.alpha.x[-1]) + 1, 1) - self.alpha.x[0] ) xticks *= (self.alpha.x.size - 1) / (self.alpha.x[-1] - self.alpha.x[0]) if len(xticks) > 5: xticks = xticks[:: len(xticks) // 5] # Add colorbar limits to the ticks (May not give ideal results) xticks = np.append(visible[0], xticks) xticks = np.append(visible[-1], xticks) # remove dupes xticks = list(set(xticks)) ax.yaxis.set_ticks(xticks) def x_format(x, pos): val = ( x * (self.alpha.x[-1] - self.alpha.x[0]) / (self.alpha.x.size - 1) + self.alpha.x[0] ) if log_scale: val = 10**val if label_fmt is None: if abs(val) < 1.0e-3 or abs(val) > 1.0e4: if not val == 0.0: e = np.floor(np.log10(abs(val))) return rf"${val / 10.0**e:.2f}\times 10^{{ {int(e):d} }}$" else: return r"$0$" else: return f"{val:.1g}" else: return label_fmt % (val) ax.yaxis.set_major_formatter(FuncFormatter(x_format)) yticks = np.linspace(0, 1, 2, endpoint=True) * max_alpha ax.xaxis.set_ticks(yticks) def y_format(y, pos): s = f"{y:0.2f}" return s ax.xaxis.set_major_formatter(FuncFormatter(y_format)) ax.set_xlim(0.0, max_alpha) ax.get_xaxis().set_ticks([]) ax.set_ylim(visible[0].item(), visible[-1].item()) ax.tick_params(axis="y", colors="white", labelsize=label_fontsize) ax.set_ylabel(label, color="white", size=size * resolution / 512.0)
[docs] def sample_colormap(self, v, w, alpha=None, colormap="gist_stern", col_bounds=None): r"""Add a Gaussian based on an existing colormap. Constructing pleasing Gaussians in a transfer function can pose some challenges, so this function will add a single Gaussian whose colors are taken from a colormap scaled between the bounds of the transfer function. As with `TransferFunction.add_gaussian`, the value is calculated as :math:`f(x) = h \exp{-(x-x_0)^2 / w}` but with the height for each color calculated from the colormap. Parameters ---------- v : float The value at which the Gaussian is to be added. w : float The relative width (:math:`w` in the above equation.) alpha : float, optional The alpha value height for the Gaussian colormap : string, optional An acceptable colormap. See either yt.visualization.color_maps or https://scipy-cookbook.readthedocs.io/items/Matplotlib_Show_colormaps.html . col_bounds: array_like, optional Limits ([min, max]) the values over which the colormap spans to these values. Useful for sampling an entire colormap over a range smaller than the transfer function bounds. See Also -------- ColorTransferFunction.add_layers : Many-at-a-time adder Examples -------- >>> tf = ColorTransferFunction((-10.0, -5.0)) >>> tf.sample_colormap(-7.0, 0.01, colormap="cmyt.arbre") """ import matplotlib as mpl v = np.float64(v) if col_bounds is None: rel = (v - self.x_bounds[0]) / (self.x_bounds[1] - self.x_bounds[0]) else: rel = (v - col_bounds[0]) / (col_bounds[1] - col_bounds[0]) cmap = mpl.colormaps[colormap] r, g, b, a = cmap(rel) if alpha is None: alpha = a self.add_gaussian(v, w, [r, g, b, alpha]) mylog.debug( "Adding gaussian at %s with width %s and colors %s", v, w, (r, g, b, alpha) )
[docs] def map_to_colormap( self, mi, ma, scale=1.0, colormap="gist_stern", scale_func=None ): r"""Map a range of values to a full colormap. Given a minimum and maximum value in the TransferFunction, map a full colormap over that range at an alpha level of `scale`. Optionally specify a scale_func function that modifies the alpha as a function of the transfer function value. Parameters ---------- mi : float The start of the TransferFunction to map the colormap ma : float The end of the TransferFunction to map the colormap scale: float, optional The alpha value to be used for the height of the transfer function. Larger values will be more opaque. colormap : string, optional An acceptable colormap. See either yt.visualization.color_maps or https://scipy-cookbook.readthedocs.io/items/Matplotlib_Show_colormaps.html . scale_func: function(:obj:`!value`, :obj:`!minval`, :obj:`!maxval`), optional A user-defined function that can be used to scale the alpha channel as a function of the TransferFunction field values. Function maps value to somewhere between minval and maxval. Examples -------- >>> def linramp(vals, minval, maxval): ... return (vals - vals.min()) / (vals.max() - vals.min()) >>> tf = ColorTransferFunction((-10.0, -5.0)) >>> tf.map_to_colormap(-8.0, -6.0, scale=10.0, colormap="cmyt.arbre") >>> tf.map_to_colormap( ... -6.0, -5.0, scale=10.0, colormap="cmyt.arbre", scale_func=linramp ... ) """ import matplotlib as mpl mi = np.float64(mi) ma = np.float64(ma) rel0 = int( self.nbins * (mi - self.x_bounds[0]) / (self.x_bounds[1] - self.x_bounds[0]) ) rel1 = int( self.nbins * (ma - self.x_bounds[0]) / (self.x_bounds[1] - self.x_bounds[0]) ) rel0 = max(rel0, 0) rel1 = min(rel1, self.nbins - 1) + 1 tomap = np.linspace(0.0, 1.0, num=rel1 - rel0) cmap = mpl.colormaps[colormap] cc = cmap(tomap) if scale_func is None: scale_mult = 1.0 else: scale_mult = scale_func(tomap, 0.0, 1.0) self.red.y[rel0:rel1] = cc[:, 0] * scale_mult self.green.y[rel0:rel1] = cc[:, 1] * scale_mult self.blue.y[rel0:rel1] = cc[:, 2] * scale_mult self.alpha.y[rel0:rel1] = scale * cc[:, 3] * scale_mult self.features.append( ( "map_to_colormap", f"start(x):{mi:3.2g}", f"stop(x):{ma:3.2g}", f"value(y):{scale:3.2g}", ) )
[docs] def add_layers( self, N, w=None, mi=None, ma=None, alpha=None, colormap="gist_stern", col_bounds=None, ): r"""Add a set of Gaussians based on an existing colormap. Constructing pleasing Gaussians in a transfer function can pose some challenges, so this function will add several evenly-spaced Gaussians whose colors are taken from a colormap scaled between the bounds of the transfer function. For each Gaussian to be added, `ColorTransferFunction.sample_colormap` is called. Parameters ---------- N : int How many Gaussians to add w : float The relative width of each Gaussian. If not included, it is calculated as 0.001 * (max_val - min_val) / N mi : float, optional If only a subset of the data range is to have the Gaussians added, this is the minimum for that subset ma : float, optional If only a subset of the data range is to have the Gaussians added, this is the maximum for that subset alpha : list of floats, optional The alpha value height for each Gaussian. If not supplied, it is set as 1.0 everywhere. colormap : string, optional An acceptable colormap. See either yt.visualization.color_maps or https://scipy-cookbook.readthedocs.io/items/Matplotlib_Show_colormaps.html . col_bounds: array_like, optional Limits ([min, max]) the values over which the colormap spans to these values. Useful for sampling an entire colormap over a range smaller than the transfer function bounds. See Also -------- ColorTransferFunction.sample_colormap : Single Gaussian adder Examples -------- >>> tf = ColorTransferFunction((-10.0, -5.0)) >>> tf.add_layers(8) """ if col_bounds is None: dist = self.x_bounds[1] - self.x_bounds[0] if mi is None: mi = self.x_bounds[0] + dist / (10.0 * N) if ma is None: ma = self.x_bounds[1] - dist / (10.0 * N) else: dist = col_bounds[1] - col_bounds[0] if mi is None: mi = col_bounds[0] + dist / (10.0 * N) if ma is None: ma = col_bounds[1] - dist / (10.0 * N) if w is None: w = 0.001 * (ma - mi) / N w = max(w, 1.0 / self.nbins) if alpha is None and self.grey_opacity: alpha = np.ones(N, dtype="float64") elif alpha is None and not self.grey_opacity: alpha = np.logspace(-3, 0, N) for v, a in zip(np.mgrid[mi : ma : N * 1j], alpha): self.sample_colormap(v, w, a, colormap=colormap, col_bounds=col_bounds)
[docs] def get_colormap_image(self, height, width): image = np.zeros((height, width, 3), dtype="uint8") hvals = np.mgrid[self.x_bounds[0] : self.x_bounds[1] : height * 1j] for i, f in enumerate(self.funcs[:3]): vals = np.interp(hvals, f.x, f.y) image[:, :, i] = (vals[:, None] * 255).astype("uint8") image = image[::-1, :, :] return image
[docs] def clear(self): for f in self.funcs: f.clear() self.features = []
def __repr__(self): disp = ( "<Color Transfer Function Object>:\n" + "x_bounds:[%3.2g, %3.2g] nbins:%i features:\n" % (self.x_bounds[0], self.x_bounds[1], self.nbins) ) for f in self.features: disp += f"\t{str(f)}\n" return disp
[docs] class ProjectionTransferFunction(MultiVariateTransferFunction): r"""A transfer function that defines a simple projection. To generate an interpolated, off-axis projection through a dataset, this transfer function should be used. It will create a very simple table that merely sums along each ray. Note that the end product will need to be scaled by the total width through which the rays were cast, a piece of information inaccessible to the transfer function. Parameters ---------- x_bounds : tuple of floats, optional If any of your values lie outside this range, they will be truncated. n_fields : int, optional How many fields we're going to project and pass through Notes ----- When you use this transfer function, you may need to explicitly disable logging of fields. """ def __init__(self, x_bounds=(-1e60, 1e60), n_fields=1): if n_fields > 3: raise NotImplementedError( f"supplied ${n_fields} but n_fields > 3 not implemented." ) MultiVariateTransferFunction.__init__(self) # Strip units off of x_bounds, if any x_bounds = [np.float64(xb) for xb in x_bounds] self.x_bounds = x_bounds self.nbins = 2 self.linear_mapping = TransferFunction(x_bounds, 2) self.linear_mapping.pass_through = 1 self.link_channels(0, [0, 1, 2]) # same emission for all rgb, default for i in range(n_fields): self.add_field_table(self.linear_mapping, i) self.link_channels(i, i) self.link_channels(n_fields, [3, 4, 5]) # this will remove absorption
[docs] class PlanckTransferFunction(MultiVariateTransferFunction): """ This sets up a planck function for multivariate emission and absorption. We assume that the emission is black body, which is then convolved with appropriate Johnson filters for *red*, *green* and *blue*. *T_bounds* and *rho_bounds* define the limits of tabulated emission and absorption functions. *anorm* is a "fudge factor" that defines the somewhat arbitrary normalization to the scattering approximation: because everything is done largely unit-free, and is really not terribly accurate anyway, feel free to adjust this to change the relative amount of reddening. Maybe in some future version this will be unitful. """ def __init__( self, T_bounds, rho_bounds, nbins=256, red="R", green="V", blue="B", anorm=1e6 ): MultiVariateTransferFunction.__init__(self) mscat = -1 from .UBVRI import johnson_filters for i, f in enumerate([red, green, blue]): jf = johnson_filters[f] tf = TransferFunction(T_bounds) tf.add_filtered_planck(jf["wavelen"], jf["trans"]) self.add_field_table(tf, 0, 1) self.link_channels(i, i) # 0 => 0, 1 => 1, 2 => 2 mscat = max(mscat, jf["Lchar"] ** -4) for i, f in enumerate([red, green, blue]): # Now we set up the scattering scat = (johnson_filters[f]["Lchar"] ** -4 / mscat) * anorm tf = TransferFunction(rho_bounds) mylog.debug("Adding: %s with relative scattering %s", f, scat) tf.y *= 0.0 tf.y += scat self.add_field_table(tf, 1, weight_field_id=1) self.link_channels(i + 3, i + 3) self._normalize() self.grey_opacity = False def _normalize(self): fmax = np.array([f.y for f in self.tables[:3]]) normal = fmax.max(axis=0) for f in self.tables[:3]: f.y = f.y / normal
if __name__ == "__main__": tf = ColorTransferFunction((-20, -5)) tf.add_gaussian(-16.0, 0.4, [0.2, 0.3, 0.1]) tf.add_gaussian(-14.0, 0.8, [0.4, 0.1, 0.2]) tf.add_gaussian(-10.0, 1.0, [0.0, 0.0, 1.0]) tf.plot("tf.png")