src/blender_maxwell/node_trees/maxwell_sim_nodes/managed_objs/managed_bl_image.py

@@ -2,10 +2,13 @@ import io
 import typing as typ
 
 import bpy
+import jax
+import jax.numpy as jnp
+import matplotlib
 import matplotlib.axis as mpl_ax
 import numpy as np
 
-from blender_maxwell.utils import image_ops, logger
+from blender_maxwell.utils import logger
 
 from .. import contracts as ct
 from . import base
@@ -15,6 +18,64 @@ log = logger.get(__name__)
 AREA_TYPE = 'IMAGE_EDITOR'
 SPACE_TYPE = 'IMAGE_EDITOR'
 
+# Colormap
+_MPL_CM = matplotlib.cm.get_cmap('viridis', 512)
+VIRIDIS_COLORMAP = jnp.array([_MPL_CM(i)[:3] for i in range(512)])
+
+
+####################
+# - Image Functions
+####################
+def apply_colormap(normalized_data, colormap):
+	# Linear interpolation between colormap points
+	n_colors = colormap.shape[0]
+	indices = normalized_data * (n_colors - 1)
+	lower_idx = jnp.floor(indices).astype(jnp.int32)
+	upper_idx = jnp.ceil(indices).astype(jnp.int32)
+	alpha = indices - lower_idx
+
+	lower_colors = jax.vmap(lambda i: colormap[i])(lower_idx)
+	upper_colors = jax.vmap(lambda i: colormap[i])(upper_idx)
+
+	return (1 - alpha)[..., None] * lower_colors + alpha[..., None] * upper_colors
+
+
+@jax.jit
+def rgba_image_from_2d_map__viridis(map_2d):
+	amplitude = jnp.abs(map_2d)
+	amplitude_normalized = (amplitude - amplitude.min()) / (
+		amplitude.max() - amplitude.min()
+	)
+	rgb_array = apply_colormap(amplitude_normalized, VIRIDIS_COLORMAP)
+	alpha_channel = jnp.ones_like(amplitude_normalized)
+	return jnp.dstack((rgb_array, alpha_channel))
+
+
+@jax.jit
+def rgba_image_from_2d_map__grayscale(map_2d):
+	amplitude = jnp.abs(map_2d)
+	amplitude_normalized = (amplitude - amplitude.min()) / (
+		amplitude.max() - amplitude.min()
+	)
+	rgb_array = jnp.stack([amplitude_normalized] * 3, axis=-1)
+	alpha_channel = jnp.ones_like(amplitude_normalized)
+	return jnp.dstack((rgb_array, alpha_channel))
+
+
+def rgba_image_from_2d_map(map_2d, colormap: str | None = None):
+	"""RGBA Image from a map of 2D coordinates to values.
+
+	Parameters:
+		map_2d: Shape (width, height, value).
+
+	Returns:
+		Image as a JAX array of shape (height, width, 4)
+	"""
+	if colormap == 'VIRIDIS':
+		return rgba_image_from_2d_map__viridis(map_2d)
+	if colormap == 'GRAYSCALE':
+		return rgba_image_from_2d_map__grayscale(map_2d)
+
 
 ####################
 # - Managed BL Image
@@ -174,7 +235,7 @@ class ManagedBLImage(base.ManagedObj):
 		self, map_2d, colormap: str | None = 'VIRIDIS', bl_select: bool = False
 	):
 		self.data_to_image(
-			lambda _: image_ops.rgba_image_from_2d_map(map_2d, colormap=colormap),
+			lambda _: rgba_image_from_2d_map(map_2d, colormap=colormap),
 			bl_select=bl_select,
 		)
 

src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/extract_data.py

@@ -1,6 +1,5 @@
 import typing as typ
 
-import enum
 import bpy
 import jax
 import jax.numpy as jnp
@@ -16,12 +15,7 @@ log = logger.get(__name__)
 
 
 class ExtractDataNode(base.MaxwellSimNode):
-	"""Node for extracting data from particular objects.
-
-	Attributes:
-		extract_filter: Identifier for data to extract from the input.
-
-	"""
+	"""Node for extracting data from particular objects."""
 
 	node_type = ct.NodeType.ExtractData
 	bl_label = 'Extract'
@@ -38,10 +32,11 @@ class ExtractDataNode(base.MaxwellSimNode):
 	####################
 	# - Properties
 	####################
-	extract_filter: enum.Enum = bl_cache.BLField(
-		None,
-		prop_ui=True,
-		enum_cb=lambda self, _: self.search_extract_filters(),
+	extract_filter: bpy.props.StringProperty(
+		name='Extract Filter',
+		description='Data to extract from the input',
+		search=lambda self, _, edit_text: self.search_extract_filters(edit_text),
+		update=lambda self, context: self.on_prop_changed('extract_filter', context),
 	)
 
 	# Sim Data
@@ -54,30 +49,41 @@ class ExtractDataNode(base.MaxwellSimNode):
 	####################
 	# - Computed Properties
 	####################
-	@property
+	@bl_cache.cached_bl_property(persist=False)
 	def has_sim_data(self) -> bool:
-		return self.active_socket_set == 'Sim Data' and self.sim_data_monitor_nametype
+		return (
+			self.active_socket_set == 'Sim Data'
+			and self.inputs['Sim Data'].is_linked
+			and self.sim_data_monitor_nametype
+		)
 
-	@property
+	@bl_cache.cached_bl_property(persist=False)
 	def has_monitor_data(self) -> bool:
-		return self.active_socket_set == 'Monitor Data' and self.monitor_data_type
+		return (
+			self.active_socket_set == 'Monitor Data'
+			and self.inputs['Monitor Data'].is_linked
+			and self.monitor_data_type
+		)
 
 	####################
 	# - Extraction Filter Search
 	####################
-	def search_extract_filters(self) -> list[ct.BLEnumElement]:
+	def search_extract_filters(self, edit_text: str) -> list[tuple[str, str, str]]:
 		if self.has_sim_data:
 			return [
-				(monitor_name, monitor_name, monitor_type.removesuffix('Data'), '', i)
-				for i, (monitor_name, monitor_type) in enumerate(
-					self.sim_data_monitor_nametype.items()
+				(
+					monitor_name,
+					monitor_type.removesuffix('Data'),
 				)
+				for monitor_name, monitor_type in self.sim_data_monitor_nametype.items()
+				if edit_text == '' or edit_text.lower() in monitor_name.lower()
 			]
 
 		if self.has_monitor_data:
 			return [
-				(component_name, component_name, f'ℂ {component_name[1]}-Pol', '', i)
-				for i, component_name in enumerate(self.monitor_data_components)
+				(component_name, f'ℂ {component_name[1]}-Pol')
+				for component_name in self.monitor_data_components
+				if (edit_text == '' or edit_text.lower() in component_name.lower())
 			]
 
 		return []
@@ -86,7 +92,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 	# - UI
 	####################
 	def draw_props(self, _: bpy.types.Context, col: bpy.types.UILayout) -> None:
-		col.prop(self, self.blfields['extract_filter'], text='')
+		col.prop(self, 'extract_filter', text='')
 
 	def draw_info(self, _: bpy.types.Context, col: bpy.types.UILayout) -> None:
 		if self.has_sim_data or self.has_monitor_data:
@@ -102,9 +108,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 			row = col.row()
 			box = row.box()
 			grid = box.grid_flow(row_major=True, columns=2, even_columns=True)
-			for name, desc in [
-				(name, desc) for idname, name, desc, *_ in self.search_extract_filters()
-			]:
+			for name, desc in self.search_extract_filters(edit_text=''):
 				grid.label(text=name)
 				grid.label(text=desc if desc else '')
 
@@ -116,7 +120,6 @@ class ExtractDataNode(base.MaxwellSimNode):
 		prop_name='active_socket_set',
 		input_sockets={'Sim Data', 'Monitor Data'},
 		input_sockets_optional={'Sim Data': True, 'Monitor Data': True},
-		run_on_init=True,
 	)
 	def on_sim_data_changed(self, input_sockets: dict):
 		if input_sockets['Sim Data'] is not None:
@@ -127,8 +130,6 @@ class ExtractDataNode(base.MaxwellSimNode):
 					'Sim Data'
 				].monitor_data.items()
 			}
-		elif self.sim_data_monitor_nametype:
-			self.sim_data_monitor_nametype = {}
 
 		if input_sockets['Monitor Data'] is not None:
 			# Monitor Data Type
@@ -167,14 +168,18 @@ class ExtractDataNode(base.MaxwellSimNode):
 					'Htheta',
 					'Hphi',
 				]
-		else:
-			if self.monitor_data_type:
-				self.monitor_data_type = ''
-			if self.monitor_data_components:
-				self.monitor_data_components = []
 
 		# Invalidate Computed Property Caches
-		self.extract_filter = bl_cache.Signal.ResetEnumItems
+		self.has_sim_data = bl_cache.Signal.InvalidateCache
+		self.has_monitor_data = bl_cache.Signal.InvalidateCache
+
+		# Reset Extraction Filter
+		## The extraction filter that was set before may not be valid anymore.
+		## If so, simply remove it.
+		if self.extract_filter not in [
+			el[0] for el in self.search_extract_filters(edit_text='')
+		]:
+			self.extract_filter = ''
 
 	####################
 	# - Output: Value
@@ -220,22 +225,16 @@ class ExtractDataNode(base.MaxwellSimNode):
 		kind=ct.FlowKind.Info,
 		props={'monitor_data_type', 'extract_filter'},
 		input_sockets={'Monitor Data'},
-		input_socket_kinds={'Monitor Data': ct.FlowKind.Value},
 	)
 	def compute_extracted_data_info(
 		self, props: dict, input_sockets: dict
-	) -> ct.InfoFlow:
-		# Retrieve XArray
-		if (
-			input_sockets['Monitor Data'] is not None
-			and props['extract_filter'] != 'NONE'
-		):
-			xarr = getattr(input_sockets['Monitor Data'], props['extract_filter'])
-		else:
+	) -> ct.InfoFlow:  # noqa: PLR0911
+		if input_sockets['Monitor Data'] is None or not props['extract_filter']:
 			return ct.InfoFlow()
 
-		# Compute InfoFlow from XArray
-		## XYZF: Field / Permittivity / FieldProjectionCartesian
+		xarr = getattr(input_sockets['Monitor Data'], props['extract_filter'])
+
+		# XYZF: Field / Permittivity / FieldProjectionCartesian
 		if props['monitor_data_type'] in {
 			'Field',
 			'Permittivity',
@@ -258,7 +257,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 				},
 			)
 
-		## XYZT: FieldTime
+		# XYZT: FieldTime
 		if props['monitor_data_type'] == 'FieldTime':
 			return ct.InfoFlow(
 				dim_names=['x', 'y', 'z', 't'],
@@ -277,7 +276,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 				},
 			)
 
-		## F: Flux
+		# F: Flux
 		if props['monitor_data_type'] == 'Flux':
 			return ct.InfoFlow(
 				dim_names=['f'],
@@ -290,7 +289,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 				},
 			)
 
-		## T: FluxTime
+		# T: FluxTime
 		if props['monitor_data_type'] == 'FluxTime':
 			return ct.InfoFlow(
 				dim_names=['t'],
@@ -303,7 +302,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 				},
 			)
 
-		## RThetaPhiF: FieldProjectionAngle
+		# RThetaPhiF: FieldProjectionAngle
 		if props['monitor_data_type'] == 'FieldProjectionAngle':
 			return ct.InfoFlow(
 				dim_names=['r', 'theta', 'phi', 'f'],
@@ -329,7 +328,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 				},
 			)
 
-		## UxUyRF: FieldProjectionKSpace
+		# UxUyRF: FieldProjectionKSpace
 		if props['monitor_data_type'] == 'FieldProjectionKSpace':
 			return ct.InfoFlow(
 				dim_names=['ux', 'uy', 'r', 'f'],
@@ -353,7 +352,7 @@ class ExtractDataNode(base.MaxwellSimNode):
 				},
 			)
 
-		## OrderxOrderyF: Diffraction
+		# OrderxOrderyF: Diffraction
 		if props['monitor_data_type'] == 'Diffraction':
 			return ct.InfoFlow(
 				dim_names=['orders_x', 'orders_y', 'f'],

src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/math/filter_math.py

@@ -43,8 +43,8 @@ class FilterMathNode(base.MaxwellSimNode):
 		prop_ui=True, enum_cb=lambda self, _: self.search_operations()
 	)
 
-	dim: enum.Enum = bl_cache.BLField(
-		None, prop_ui=True, enum_cb=lambda self, _: self.search_dims()
+	dim: str = bl_cache.BLField(
+		'', prop_ui=True, str_cb=lambda self, _, edit_text: self.search_dims(edit_text)
 	)
 
 	dim_names: list[str] = bl_cache.BLField([])
@@ -64,24 +64,22 @@ class FilterMathNode(base.MaxwellSimNode):
 				('FIX', 'del a | i≈v', 'Fix Coordinate'),
 			]
 
-		return [(*item, '', i) for i, item in enumerate(items)]
+		return items
 
 	####################
 	# - Dim Search
 	####################
-	def search_dims(self) -> list[ct.BLEnumElement]:
+	def search_dims(self, edit_text: str) -> list[tuple[str, str, str]]:
 		if self.dim_names:
 			dims = [
-				(dim_name, dim_name, dim_name, '', i)
-				for i, dim_name in enumerate(self.dim_names)
+				(dim_name, dim_name)
+				for dim_name in self.dim_names
+				if edit_text == '' or edit_text.lower() in dim_name.lower()
 			]
 
 			# Squeeze: Dimension Must Have Length=1
-			## We must also correct the "NUMBER" of the enum.
 			if self.operation == 'SQUEEZE':
-				filtered_dims = [dim for dim in dims if self.dim_lens[dim[0]] == 1]
-				return [(*dim[:-1], i) for i, dim in enumerate(filtered_dims)]
-
+				return [dim for dim in dims if self.dim_lens[dim[0]] == 1]
 			return dims
 		return []
 
@@ -109,7 +107,6 @@ class FilterMathNode(base.MaxwellSimNode):
 		input_sockets={'Data'},
 		input_socket_kinds={'Data': ct.FlowKind.Info},
 		input_sockets_optional={'Data': True},
-		run_on_init=True,
 	)
 	def on_any_change(self, props: dict, input_sockets: dict):
 		# Set Dimension Names from InfoFlow
@@ -123,8 +120,8 @@ class FilterMathNode(base.MaxwellSimNode):
 			self.dim_names = []
 			self.dim_lens = {}
 
-		# Reset Enum
-		self.dim = bl_cache.Signal.ResetEnumItems
+		# Reset String Searcher
+		self.dim = bl_cache.Signal.ResetStrSearch
 
 	@events.on_value_changed(
 		prop_name='dim',
@@ -135,7 +132,10 @@ class FilterMathNode(base.MaxwellSimNode):
 	)
 	def on_dim_change(self, props: dict, input_sockets: dict):
 		# Add/Remove Input Socket "Value"
-		if props['active_socket_set'] == 'By Dim Value' and props['dim'] != 'NONE':
+		if (
+			props['active_socket_set'] == 'By Dim Value'
+			and props['dim'] in input_sockets['Data'].dim_names
+		):
 			# Get Current and Wanted Socket Defs
 			current_socket_def = self.loose_input_sockets.get('Value')
 			wanted_socket_def = sockets.SOCKET_DEFS[
@@ -167,7 +167,7 @@ class FilterMathNode(base.MaxwellSimNode):
 
 		# Compute Bound/Free Parameters
 		func_args = [int] if props['active_socket_set'] == 'By Dim Value' else []
-		if props['dim'] != 'NONE':
+		if props['dim']:
 			axis = info.dim_names.index(props['dim'])
 		else:
 			msg = 'Dimension cannot be empty'
@@ -225,7 +225,7 @@ class FilterMathNode(base.MaxwellSimNode):
 
 		# Compute Bound/Free Parameters
 		## Empty Dimension -> Empty InfoFlow
-		if props['dim'] != 'NONE':
+		if props['dim']:
 			axis = info.dim_names.index(props['dim'])
 		else:
 			return ct.InfoFlow()
@@ -272,7 +272,7 @@ class FilterMathNode(base.MaxwellSimNode):
 			in [
 				('By Dim Value', 'FIX'),
 			]
-			and props['dim'] != 'NONE'
+			and props['dim']
 			and input_sockets['Value'] is not None
 		):
 			# Compute IDX Corresponding to Coordinate Value

src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/viz.py

@@ -1,9 +1,8 @@
-import enum
 import typing as typ
 
 import bpy
 
-from blender_maxwell.utils import bl_cache, image_ops, logger
+from blender_maxwell.utils import logger
 
 from ... import contracts as ct
 from ... import managed_objs, sockets
@@ -13,12 +12,7 @@ log = logger.get(__name__)
 
 
 class VizNode(base.MaxwellSimNode):
-	"""Node for visualizing simulation data, by querying its monitors.
-
-	Attributes:
-		colormap: Colormap to apply to 0..1 output.
-
-	"""
+	"""Node for visualizing simulation data, by querying its monitors."""
 
 	node_type = ct.NodeType.Viz
 	bl_label = 'Viz'
@@ -40,25 +34,31 @@ class VizNode(base.MaxwellSimNode):
 	#####################
 	## - Properties
 	#####################
-	colormap: image_ops.Colormap = bl_cache.BLField(
-		image_ops.Colormap.Viridis, prop_ui=True
+	colormap: bpy.props.EnumProperty(
+		name='Colormap',
+		description='Colormap to apply to grayscale output',
+		items=[
+			('VIRIDIS', 'Viridis', 'Good default colormap'),
+			('GRAYSCALE', 'Grayscale', 'Barebones'),
+		],
+		default='VIRIDIS',
+		update=lambda self, context: self.on_prop_changed('colormap', context),
 	)
 
 	#####################
 	## - UI
 	#####################
 	def draw_props(self, _: bpy.types.Context, col: bpy.types.UILayout):
-		col.prop(self, self.blfields['colormap'], text='')
+		col.prop(self, 'colormap')
 
 	#####################
 	## - Plotting
 	#####################
 	@events.on_show_plot(
 		managed_objs={'plot'},
-		props={'colormap'},
 		input_sockets={'Data'},
 		input_socket_kinds={'Data': ct.FlowKind.Array},
-		input_sockets_optional={'Data': True},
+		props={'colormap'},
 		stop_propagation=True,
 	)
 	def on_show_plot(
@@ -67,7 +67,6 @@ class VizNode(base.MaxwellSimNode):
 		input_sockets: dict,
 		props: dict,
 	):
-		if input_sockets['Data'] is not None:
 		managed_objs['plot'].map_2d_to_image(
 			input_sockets['Data'].values,
 			colormap=props['colormap'],

src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/base.py

@@ -1034,16 +1034,6 @@ class MaxwellSimNode(bpy.types.Node):
 		## Blender will automatically add .001 so that `self.name` is unique.
 		self.sim_node_name = self.name
 
-		# Event Methods
-		## Run any 'DataChanged' methods with 'run_on_init' set.
-		## -> Copying a node _arguably_ re-initializes the new node.
-		for event_method in [
-			event_method
-			for event_method in self.event_methods_by_event[ct.FlowEvent.DataChanged]
-			if event_method.callback_info.run_on_init
-		]:
-			event_method(self)
-
 	def free(self) -> None:
 		"""Cleans various instance-associated data up, so the node can be cleanly deleted.
 

src/blender_maxwell/utils/bl_cache.py

@@ -740,7 +740,7 @@ class BLField:
 						str(value),
 						AttrType.to_name(value),
 						AttrType.to_name(value),  ## TODO: From AttrType.__doc__
-						AttrType.to_icon(value),
+						AttrType.to_icon(),
 						i if not self._enum_many else 2**i,
 					)
 					for i, value in enumerate(list(AttrType))
@@ -824,12 +824,6 @@ class BLField:
 
 	def __set__(self, bl_instance: BLInstance | None, value: typ.Any) -> None:
 		if value == Signal.ResetEnumItems:
-			old_items = self._safe_enum_cb(bl_instance, None)
-			current_items = self._enum_cb(bl_instance, None)
-
-			# Only Change if Changes Need Making
-			## <Proverb. 'Fun things to say in jail'. Ca 887BCE. >
-			if old_items != current_items:
 			# Set Enum to First Item
 			## Prevents the seemingly "missing" enum element bug.
 			## -> Caused by the old int still trying to hang on after.
@@ -837,7 +831,7 @@ class BLField:
 			## -> Infinite recursion if we don't check current value.
 			## -> May cause a hiccup (chains will trigger twice)
 			## To work, there **must** be a guaranteed-available string at 0,0.
-				first_old_value = old_items(bl_instance, None)[0][0]
+			first_old_value = self._safe_enum_cb(bl_instance, None)[0][0]
 			current_value = self._cached_bl_property.__get__(
 				bl_instance, bl_instance.__class__
 			)
@@ -853,11 +847,6 @@ class BLField:
 			self._cached_bl_property.__set__(bl_instance, Signal.InvalidateCache)
 
 		elif value == Signal.ResetStrSearch:
-			old_items = self._safe_str_cb(bl_instance, None)
-			current_items = self._str_cb(bl_instance, None)
-
-			# Only Change if Changes Need Making
-			if old_items != current_items:
 			# Set String to ''
 			## Prevents the presence of an invalid value not in the new search.
 			## -> Infinite recursion if we don't check current value for ''.

src/blender_maxwell/utils/image_ops.py

@@ -1,108 +0,0 @@
-"""Useful image processing operations for use in the addon."""
-
-import enum
-import typing as typ
-
-import jax
-import jax.numpy as jnp
-import jaxtyping as jtyp
-import matplotlib
-
-from blender_maxwell import contracts as ct
-from blender_maxwell.utils import logger
-
-log = logger.get(__name__)
-
-####################
-# - Constants
-####################
-_MPL_CM = matplotlib.cm.get_cmap('viridis', 512)
-VIRIDIS_COLORMAP: jtyp.Float32[jtyp.Array, '512 3'] = jnp.array(
-	[_MPL_CM(i)[:3] for i in range(512)]
-)
-
-
-class Colormap(enum.StrEnum):
-	"""Available colormaps.
-
-	Attributes:
-		Viridis: Good general-purpose colormap.
-		Grayscale: Simple black and white mapping.
-	"""
-
-	Viridis = enum.auto()
-	Grayscale = enum.auto()
-
-	@staticmethod
-	def to_name(value: typ.Self) -> str:
-		return {
-			Colormap.Viridis: 'Viridis',
-			Colormap.Grayscale: 'Grayscale',
-		}[value]
-
-	@staticmethod
-	def to_icon(value: typ.Self) -> ct.BLIcon:
-		return ''
-
-
-####################
-# - Colormap: (X,Y,1 -> Value) -> (X,Y,4 -> Value)
-####################
-def apply_colormap(
-	normalized_data: jtyp.Float32[jtyp.Array, 'width height 4'],
-	colormap: jtyp.Float32[jtyp.Array, '512 3'],
-):
-	# Linear interpolation between colormap points
-	n_colors = colormap.shape[0]
-	indices = normalized_data * (n_colors - 1)
-	lower_idx = jnp.floor(indices).astype(jnp.int32)
-	upper_idx = jnp.ceil(indices).astype(jnp.int32)
-	alpha = indices - lower_idx
-
-	lower_colors = jax.vmap(lambda i: colormap[i])(lower_idx)
-	upper_colors = jax.vmap(lambda i: colormap[i])(upper_idx)
-
-	return (1 - alpha)[..., None] * lower_colors + alpha[..., None] * upper_colors
-
-
-@jax.jit
-def rgba_image_from_2d_map__viridis(map_2d: jtyp.Float32[jtyp.Array, 'width height 4']):
-	amplitude = jnp.abs(map_2d)
-	amplitude_normalized = (amplitude - amplitude.min()) / (
-		amplitude.max() - amplitude.min()
-	)
-	rgb_array = apply_colormap(amplitude_normalized, VIRIDIS_COLORMAP)
-	alpha_channel = jnp.ones_like(amplitude_normalized)
-	return jnp.dstack((rgb_array, alpha_channel))
-
-
-@jax.jit
-def rgba_image_from_2d_map__grayscale(
-	map_2d: jtyp.Float32[jtyp.Array, 'width height 4'],
-):
-	amplitude = jnp.abs(map_2d)
-	amplitude_normalized = (amplitude - amplitude.min()) / (
-		amplitude.max() - amplitude.min()
-	)
-	rgb_array = jnp.stack([amplitude_normalized] * 3, axis=-1)
-	alpha_channel = jnp.ones_like(amplitude_normalized)
-	return jnp.dstack((rgb_array, alpha_channel))
-
-
-def rgba_image_from_2d_map(
-	map_2d: jtyp.Float32[jtyp.Array, 'width height 4'], colormap: str | None = None
-):
-	"""RGBA Image from a map of 2D coordinates to values.
-
-	Parameters:
-		map_2d: The 2D value map.
-
-	Returns:
-		Image as a JAX array of shape (height, width, 4)
-	"""
-	if colormap == Colormap.Viridis:
-		return rgba_image_from_2d_map__viridis(map_2d)
-	if colormap == Colormap.Grayscale:
-		return rgba_image_from_2d_map__grayscale(map_2d)
-
-	return rgba_image_from_2d_map__grayscale(map_2d)