"""Defines render types for setting up a RenderOutput object"""
import abc
import colorsys
from enum import Enum
from typing import Annotated, Literal
import pydantic as pd
import unyt as u
from flow360_schema.framework.base_model import Flow360BaseModel
from flow360_schema.framework.entity.geometric_types import Axis, Vector
from flow360_schema.framework.expression import (
Expression,
UnytQuantity,
UserVariable,
ValueOrExpression,
get_input_value_dimensions,
get_input_value_length,
solver_variable_to_user_variable,
)
from flow360_schema.framework.expression.utils import is_runtime_expression
from flow360_schema.framework.physical_dimensions import Angle, Length
from flow360_schema.models.entities.output_entities import (
_infer_units_by_unit_system as infer_units_by_unit_system,
)
from flow360_schema.models.entities.output_entities import (
_should_skip_unit_system_inference as is_variable_with_unit_system_as_units,
)
from flow360_schema.models.simulation.outputs.output_fields import CommonFieldNames
# Re-exported so the public `render_config.Camera` / `render_config.Viewpoint` access
# keeps working; the canonical definition lives in simulation/camera.py.
from flow360_schema.models.simulation.camera import Camera as Camera
from flow360_schema.models.simulation.camera import Viewpoint as Viewpoint
Int8 = Annotated[int, pd.Field(ge=0, le=255)]
Color = tuple[Int8, Int8, Int8]
Alpha = Annotated[float, pd.Field(ge=0.0, le=1.0)]
# A colormap key is an RGB triple, optionally with a 4th alpha channel (0-1). The
# alpha defines the per-key opacity transfer function used by volume rendering;
# keys authored without alpha fall back to a value-proportional ramp (the
# historical behaviour), so existing colormaps render unchanged.
ColormapKey = tuple[Int8, Int8, Int8] | tuple[Int8, Int8, Int8, Alpha]
class Resolution(Flow360BaseModel):
"""
:class:`Resolution` defines the output image resolution, in pixels, for a render.
Example
-------
>>> Resolution(width=1920, height=1080)
"""
width: pd.PositiveInt = pd.Field(1920, description="Output image width in pixels")
height: pd.PositiveInt = pd.Field(1080, description="Output image height in pixels")
class AmbientLight(Flow360BaseModel):
"""
:class:`AmbientLight` controls uniform ambient lighting in the scene.
Example
-------
>>> AmbientLight(
... intensity=0.4,
... color=(255, 255, 255)
... )
"""
type_name: Literal["AmbientLight"] = pd.Field("AmbientLight", frozen=True)
intensity: float = pd.Field(ge=0, description="Light intensity multiplier")
color: Color = pd.Field(description="Color of the ambient light")
class DirectionalLight(Flow360BaseModel):
"""
:class:`DirectionalLight` defines a directional light source with intensity and color.
Example
-------
>>> DirectionalLight(
... intensity=1.0,
... color=(255, 255, 255),
... direction=(-1, -1, -1),
... )
"""
type_name: Literal["DirectionalLight"] = pd.Field("DirectionalLight", frozen=True)
intensity: float = pd.Field(ge=0, description="Light intensity multiplier")
color: Color = pd.Field(description="Color of the directional light beam")
direction: Axis = pd.Field(description="The direction of the light beam (all beams are parallel)")
class Lighting(Flow360BaseModel):
"""
:class:`Lighting` defines ambient and directional lighting for rendering.
Example
-------
>>> Lighting.default()
"""
type_name: Literal["Lighting"] = pd.Field("Lighting", frozen=True)
directional: DirectionalLight = pd.Field(
description="Directional component of the light (falls from a single direction)"
)
ambient: AmbientLight | None = pd.Field(
description="Ambient component of the light (applied from all directions equally)"
)
@classmethod
def default(cls, direction=(-1.0, -1.0, -1.0)):
"""
Returns the default lighting configuration.
Example
-------
>>> light = Lighting.default()
"""
return Lighting(
ambient=AmbientLight(intensity=0.4, color=(255, 255, 255)),
directional=DirectionalLight(intensity=1.0, color=(255, 255, 255), direction=direction),
)
class BackgroundBase(Flow360BaseModel, metaclass=abc.ABCMeta):
"""
:class:`RenderBackgroundBase` is an abstract base class for all background types.
"""
type_name: str = pd.Field(default="", frozen=True)
class SolidBackground(BackgroundBase):
"""
:class:`SolidBackground` defines a single-color background.
Example
-------
>>> SolidBackground(color=(200, 200, 255))
"""
type_name: Literal["SolidBackground"] = pd.Field("SolidBackground", frozen=True)
color: Color = pd.Field(description="Flat background color")
class SkyboxTexture(str, Enum):
"""
:class:`SkyboxTexture` specifies available skybox texture presets.
Example
-------
>>> SkyboxTexture.SKY.value
'sky'
"""
SKY = "sky"
GRADIENT = "gradient"
class SkyboxBackground(BackgroundBase):
"""
:class:`SkyboxBackground` defines a skybox background using a sky or gradient texture.
Example
-------
>>> SkyboxBackground(texture=SkyboxTexture.SKY)
"""
type_name: Literal["SkyboxBackground"] = pd.Field("SkyboxBackground", frozen=True)
texture: SkyboxTexture = pd.Field(SkyboxTexture.SKY, description="Cubemap texture applied to the skybox")
class Environment(Flow360BaseModel):
"""
:class:`Environment` configures the background environment for rendering.
Example
-------
>>> Environment.simple()
"""
type_name: Literal["Environment"] = pd.Field("Environment", frozen=True)
background: SolidBackground | SkyboxBackground = pd.Field(
discriminator="type_name", description="Background image, solid or textured"
)
@classmethod
def simple(cls):
"""
Create a render environment with a solid background.
Example
-------
>>> Environment.simple()
"""
return Environment(background=SolidBackground(color=(207, 226, 230)))
@classmethod
def sky(cls):
"""
Create a render environment using a sky texture.
Example
-------
>>> Environment.sky()
"""
return Environment(background=SkyboxBackground(texture=SkyboxTexture.SKY))
@classmethod
def gradient(cls):
"""
Create a render environment using a gradient skybox.
Example
-------
>>> Environment.gradient()
"""
return Environment(background=SkyboxBackground(texture=SkyboxTexture.GRADIENT))
class MaterialBase(Flow360BaseModel, metaclass=abc.ABCMeta):
"""
:class:`MaterialBase` is an abstract base class for material definitions used during rendering.
"""
type_name: str = pd.Field("", frozen=True)
[docs]
class PBRMaterial(MaterialBase):
"""
:class:`PBRMaterial` defines a physically based rendering (PBR) material.
Example
-------
>>> PBRMaterial(color=(180, 180, 255), roughness=0.3)
"""
type_name: Literal["PBRMaterial"] = pd.Field("PBRMaterial", frozen=True)
color: Color = pd.Field(default=[255, 255, 255], description="Basic diffuse color of the material (base color)")
opacity: float = pd.Field(
default=1,
ge=0,
le=1,
description="The transparency of the material 1 is fully opaque, 0 is fully transparent",
)
roughness: float = pd.Field(
default=0.5,
ge=0,
le=1,
description="Material roughness, controls the fuzziness of reflections",
)
f0: Vector = pd.Field(
default=(0.03, 0.03, 0.03),
description="Fresnel reflection coeff. at 0 incidence angle, controls reflectivity",
)
[docs]
@classmethod
def plastic(cls, shine=0.5, opacity=1.0):
"""
Create a plastic PBR material.
Example
-------
>>> PBRMaterial.plastic(shine=0.2)
"""
return PBRMaterial(color=(255, 255, 255), opacity=opacity, roughness=1 - shine, f0=(0.03, 0.03, 0.03))
[docs]
class FieldMaterial(MaterialBase):
"""
:class:`FieldMaterial` maps scalar field values to colors for flow visualization.
Example
-------
>>> FieldMaterial.rainbow(field="pressure", min_value=0, max_value=100000)
"""
type_name: Literal["FieldMaterial"] = pd.Field("FieldMaterial", frozen=True)
opacity: float = pd.Field(
default=1,
ge=0,
le=1,
description="The transparency of the material 1 is fully opaque, 0 is fully transparent",
)
output_field: CommonFieldNames | str | UserVariable = pd.Field(
description="Scalar field applied to the surface via the colormap"
)
min: ValueOrExpression[UnytQuantity | float] = pd.Field(
description="Reference min value (in solver units) representing the left boundary of the colormap"
)
max: ValueOrExpression[UnytQuantity | float] = pd.Field(
description="Reference max value (in solver units) representing the right boundary of the colormap"
)
colormap: list[ColormapKey] = pd.Field(
description=(
"Key colors distributed evenly across the gradient (value->color mapping). "
"Each key is [r, g, b] or [r, g, b, alpha]; the optional alpha (0-1) sets the "
"per-key opacity transfer function for volume rendering. Keys without alpha use "
"a value-proportional opacity ramp."
)
)
@pd.field_validator("output_field", mode="before")
@classmethod
def _preprocess_expression_and_solver_variable(cls, value):
if isinstance(value, Expression):
raise ValueError(
f"Expression ({value}) cannot be directly used as output field, " "please define a UserVariable first."
)
return solver_variable_to_user_variable(value)
[docs]
@pd.field_validator("output_field", mode="after")
@classmethod
def check_runtime_expression(cls, v):
"""Ensure the output field is a runtime expression but not a constant value."""
if isinstance(v, UserVariable):
if not isinstance(v.value, Expression):
raise ValueError(f"The output field ({v}) cannot be a constant value.")
try:
result = v.value.evaluate(raise_on_non_evaluable=False, force_evaluate=True)
except Exception as err:
raise ValueError(f"expression evaluation failed for the output field: {err}") from err
if not is_runtime_expression(result):
raise ValueError(f"The output field ({v}) cannot be a constant value.")
return v
@pd.field_validator("min", "max", mode="before")
@classmethod
def _preprocess_range_with_unit_system(cls, value, info: pd.ValidationInfo):
if is_variable_with_unit_system_as_units(value):
return value
if info.data.get("field") is None:
# `field` validation failed.
raise ValueError("The output field is invalid and therefore unit inference is not possible.")
units = value["units"]
field = info.data["field"]
value_dimensions = get_input_value_dimensions(value=field)
value = infer_units_by_unit_system(value=value, value_dimensions=value_dimensions, unit_system=units)
return value
[docs]
@pd.field_validator("min", "max", mode="after")
@classmethod
def check_range_single_value(cls, v):
"""Ensure the min/max range is a single value."""
if get_input_value_length(v) == 0:
return v
raise ValueError(f"The min/max range ({v}) must be a scalar.")
[docs]
@pd.field_validator("min", "max", mode="after")
@classmethod
def check_range_dimensions(cls, v, info: pd.ValidationInfo):
"""Ensure the min/max range has the same dimensions as the field."""
field = info.data.get("output_field", None)
if not isinstance(field, UserVariable):
return v
range_dimensions = get_input_value_dimensions(value=v)
if range_dimensions is None:
return v
field_dimensions = get_input_value_dimensions(value=field)
if field_dimensions != range_dimensions:
raise ValueError(
f"The min/max range ({v}, dimensions:{range_dimensions}) should have the same dimensions as "
f"the output field ({field}, dimensions: {field_dimensions})."
)
return v
[docs]
@pd.field_validator("min", "max", mode="after")
@classmethod
def check_iso_value_for_string_field(cls, v, info: pd.ValidationInfo):
"""Ensure the iso_value is float when string field is used."""
field = info.data.get("output_field", None)
if isinstance(field, str) and not isinstance(v, float):
raise ValueError(
f"The output field ({field}) specified by string "
"can only be used with a nondimensional min/max range."
)
return v
[docs]
@classmethod
def rainbow(cls, field, min_value, max_value, opacity=1):
"""
Create a rainbow-style colormap for scalar fields.
Example
-------
>>> FieldMaterial.rainbow("velocity_magnitude")
"""
def _rainbow_rgb(t):
h = (((((1 - t) * 2) / 3) % 1) + 1) % 1
r, g, b = colorsys.hsv_to_rgb(h, 1.0, 1.0)
return (int(round(r * 255)), int(round(g * 255)), int(round(b * 255)))
colormap = []
for i in range(20):
t = i / (20 - 1)
colormap.append(_rainbow_rgb(t))
# Approximated from TS rainbowGradient sampling
return FieldMaterial(opacity=opacity, output_field=field, min=min_value, max=max_value, colormap=colormap)
[docs]
@classmethod
def orizon(cls, field, min_value, max_value, opacity=1):
"""
Create an Orizon-style (blue–orange) colormap.
Example
-------
>>> FieldMaterial.orizon("temperature")
"""
def _orizon_rgb(t):
h = 0.7 * t + 0.025
r, g, b = colorsys.hsv_to_rgb(h % 1.0, 0.9, 1.0)
return (int(round(r * 255)), int(round(g * 255)), int(round(b * 255)))
colormap = []
for i in range(20):
t = i / (20 - 1)
colormap.append(_orizon_rgb(t))
# Approximated from TS orizonGradient sampling
return FieldMaterial(opacity=opacity, output_field=field, min=min_value, max=max_value, colormap=colormap)
[docs]
@classmethod
def viridis(cls, field, min_value, max_value, opacity=1):
"""
Create a Viridis colormap.
Example
-------
>>> FieldMaterial.viridis("vorticity")
"""
return FieldMaterial(
opacity=opacity,
output_field=field,
min=min_value,
max=max_value,
colormap=[
(68, 1, 84),
(65, 68, 135),
(42, 120, 142),
(34, 168, 132),
(122, 209, 81),
(253, 231, 37),
],
)
[docs]
@classmethod
def magma(cls, field, min_value, max_value, opacity=1):
"""
Create a Magma colormap.
Example
-------
>>> FieldMaterial.magma("density")
"""
return FieldMaterial(
opacity=opacity,
output_field=field,
min=min_value,
max=max_value,
colormap=[
(0, 0, 4),
(86, 20, 125),
(192, 58, 118),
(253, 154, 106),
(252, 253, 191),
],
)
[docs]
@classmethod
def airflow(cls, field, min_value, max_value, opacity=1):
"""
Create an Airflow-style visualization colormap.
Example
-------
>>> FieldMaterial.airflow("pressure_coefficient")
"""
return FieldMaterial(
opacity=opacity,
output_field=field,
min=min_value,
max=max_value,
colormap=[
(0, 100, 60),
(97, 178, 156),
(123, 189, 240),
(241, 241, 240),
(254, 216, 139),
(247, 139, 141),
(252, 122, 76),
(176, 90, 249),
],
)
class SceneTransform(Flow360BaseModel):
"""
:class:`SceneTransform` applies translation, rotation, and scaling to renderable objects.
This may be
Example
-------
>>> SceneTransform(
... translation=(1, 0, 0) * u.m,
... rotation=(0, 0, 90) * u.deg,
... scale=(1, 2, 1),
... )
"""
type_name: Literal["SceneTransform"] = pd.Field("SceneTransform", frozen=True)
translation: Length.Vector3 = pd.Field((0, 0, 0) * u.m, description="Translation applied to all scene objects")
rotation: Angle.Vector3 = pd.Field(
(0, 0, 0) * u.deg, description="Rotation applied to all scene objects (Euler XYZ)"
)
scale: Vector = pd.Field((1, 1, 1), description="Scaling applied to all scene objects")
