BUG/ENH: multivariate transfer function support

[chrishavlin]

The transfer function module has a number of transfer function classes for multivariate transfer functions (a base MultiVariateTransferFunction and a PlanckTransferFunction). They build OK, but it does not seem possible to use them correctly at present.

See #2804 for a full discussion, one thing that I'll put out here though is a link to a bug report when trying to use a multivariate TF: https://docs.google.com/document/d/1SmxjoOKJYPBP4erAknK2-2IbwxdgIiCNtf_B-66yzQE/edit?tab=t.0

[chrishavlin]

Getting something! taking bits from the bug report, bits from @matthewturk 's comment here, and simply adding the MultiVariateTransferFunction to the validator:

--- a/yt/visualization/volume_rendering/render_source.py
+++ b/yt/visualization/volume_rendering/render_source.py
@@ -28,6 +28,7 @@ from .transfer_functions import (
     ColorTransferFunction,
     ProjectionTransferFunction,
     TransferFunction,
+    MultiVariateTransferFunction,
 )
 from .utils import (
     data_source_or_all,
@@ -285,6 +286,7 @@ class VolumeSource(RenderSource, abc.ABC):
             TransferFunction,
             ColorTransferFunction,
             ProjectionTransferFunction,
+            MultiVariateTransferFunction,
             type(None),
         )
         if not isinstance(value, valid_types):

I'm getting an image.

Click to see code

from yt.visualization.volume_rendering.transfer_functions import ColorTransferFunction, TransferFunction, MultiVariateTransferFunction
import yt
import numpy as np

ds = yt.load_sample("Enzo_64")

# 1. Create a transfer function for COLOR, based on Temperature.
log_temp_bounds = (3, 8)
log_dens_bounds = (-31, -27)
tf_temp = ColorTransferFunction(log_temp_bounds)
tf_temp.add_layers(10, colormap="gray", w=0.5) #Changing the colormap does not affect the render

# 2. Create a separate, simple transfer function for OPACITY, based on Density.
tf_density = TransferFunction(log_dens_bounds)
# This ramp uses a non-linear curve to make faint features visible and is amplified.
normalized_ramp = np.linspace(0.0, 1.0, tf_density.nbins)
curved_ramp = normalized_ramp**0.5
tf_density.y = 20.0 * curved_ramp


mvtf = MultiVariateTransferFunction()

# add field tables
mvtf.add_field_table(tf_temp.red, field_id=0, weight_field_id=-1, weight_table_id=-1)  # temperature → RGB
mvtf.add_field_table(tf_temp.green, field_id=0, weight_field_id=-1, weight_table_id=-1)  # temperature → RGB
mvtf.add_field_table(tf_temp.blue, field_id=0, weight_field_id=-1, weight_table_id=-1)  # temperature → RGB
mvtf.add_field_table(tf_density, field_id=1, weight_field_id=-1, weight_table_id=-1)   # density → Alpha

# Link channels
mvtf.link_channels(table_id = 0, channels = [0])
mvtf.link_channels(table_id = 1, channels = [1])
mvtf.link_channels(table_id = 2, channels = [2])
mvtf.link_channels(table_id = 3, channels = [3, 4, 5])


fields_to_render = [('gas', 'temperature'), ('gas', 'density')]
sc = yt.create_scene(ds.all_data(), fields_to_render)
source = sc[0]

# important!!! set the tf on base source, not the tfh! only works with the above diff!
source.set_transfer_function(mvtf)  

sc.save('mv_test.png')

[chrishavlin]

Ok, so spent a while yesterday trying to get a MV TF to really work and found some things. Some of these I would consider bugs (maybe), some are probably user error/preferences.

The things:

1. work in rescaled linear space to avoid nan insertion

So when applying a log to the field, values get read from the current grid object and stored as follows:

yt/yt/utilities/amr_kdtree/amr_kdtree.py

Lines 343 to 355 in f77b7c9

	vcd = grid.get_vertex_centered_data(
	self.fields, smoothed=True, no_ghost=self.no_ghost
	)
	for i, field in enumerate(self.fields):
	if self.log_fields[i]:
	v = vcd[field].astype("float64")
	v[v <= 0] = np.nan
	dds.append(np.log10(v))
	else:
	dds.append(vcd[field].astype("float64"))
	self.current_saved_grids.append(grid)
	self.current_vcds.append(dds)

I was worried that those nan values would be accumulated during volume rendering (there is no nan handling in the transfer functions that i can see), and is likely what results in having nan values in the final image arrays for some rendering. So I wrote some derived fields to rescale temperature and density between (0, 1) in log space and then used those while rendering with linear transfer functions.

I don't think I'd call this a bug... but maybe having an option to set the fill value would make sense for volume rendering? e.g., setting it to some non-nan value that falls outside the bounds of the transfer function should not affect images and would prevent nans from polluting the accumulation along a ray (or, we add nan handling during the FIT evaluation?).

2. the weighting field does not have its own log control (maybe a bug? at least confusing...)

The spot where fields are set and whether to log a field is stored uses the log value for the primary field for the weighting field:

yt/yt/visualization/volume_rendering/render_source.py

Lines 121 to 137 in f77b7c9

	def validate_volume(f):
	@wraps(f)
	def wrapper(*args, **kwargs):
	obj = args[0]
	fields = [obj.field]
	log_fields = [obj.log_field]
	if obj.weight_field is not None:
	fields.append(obj.weight_field)
	log_fields.append(obj.log_field)
	if not obj._volume_valid:
	obj.volume.set_fields(
	fields, log_fields, no_ghost=(not obj.use_ghost_zones)
	)
	obj._volume_valid = True
	return f(*args, **kwargs)
	return wrapper

This bit in particular:

         fields = [obj.field] 
         log_fields = [obj.log_field] 
         if obj.weight_field is not None: 
             fields.append(obj.weight_field) 
             log_fields.append(obj.log_field) 

IMO there should be a way to set the log_fields value for the weight field independent from the primary field. (I've got a branch where I've added a log_weight_field attribute, it is much clearer IMO).

3. an example

Here's an example notebook using an experimental branch. Note I was messing with the field rescaling a bit to see how it affects the images.