from __future__ import annotations
import typing as t
from abc import ABC, abstractmethod
import attrs
import mitsuba as mi
import numpy as np
import pint
import pinttr
from .shapes import CuboidShape, RectangleShape, Shape, SphereShape
from ..attrs import documented, parse_docs
from ..constants import EARTH_RADIUS
from ..kernel import map_cube, map_unit_cube
from ..radprops import ZGrid
from ..units import unit_context_config as ucc
from ..units import unit_context_kernel as uck
from ..units import unit_registry as ureg
[docs]
@parse_docs
@attrs.define
class SceneGeometry(ABC):
"""
Abstract base class defining a scene geometry.
Warnings
--------
If a ``zgrid`` value is passed to the constructor (instead of letting
the constructor set it automatically), its extent must be
[``ground_altitude``, ``toa_altitude``]. The constructor will raise
a :class:`ValueError` otherwise.
"""
toa_altitude: pint.Quantity = documented(
pinttr.field(default=120.0 * ureg.km, units=ucc.deferred("length")),
doc="Top-of-atmosphere level altitude. "
'Unit-enabled field (default: ``ucc["length"]``).',
default="120 km",
type="pint.Quantity",
init_type="float or quantity",
)
ground_altitude: pint.Quantity = documented(
pinttr.field(default=0.0 * ureg.km, units=ucc.deferred("length")),
doc="Baseline ground altitude. "
'Unit-enabled field (default: ``ucc["length"]``).',
default="0 km",
type="pint.Quantity",
init_type="float or quantity",
)
zgrid: ZGrid = documented(
attrs.field(
default=None,
converter=attrs.converters.optional(
lambda x: ZGrid(x) if not isinstance(x, ZGrid) else x
),
validator=attrs.validators.optional(attrs.validators.instance_of(ZGrid)),
),
doc="The altitude mesh on which the radiative properties of "
"heterogeneous atmosphere components are evaluated. "
"If unset, a default grid with one layer per 100 m (or 10 layers if "
"the atmosphere object height is less than 100 m) is used.",
type=".ZGrid",
init_type=".ZGrid, quantity or ndarray, optional",
)
def __attrs_post_init__(self) -> None:
# Set altitude grid
if self.zgrid is None:
bottom = self.ground_altitude.m_as(ureg.m)
top = self.toa_altitude.m_as(ureg.m)
step = min(100.0, (top - bottom) / 10.0)
self.zgrid = ZGrid(
ureg.convert(
np.arange(bottom, top + step * 0.1, step),
ureg.m,
ucc.get("length"),
)
)
else:
grid_bottom = self.zgrid.levels[0]
if not np.isclose(grid_bottom, self.ground_altitude):
raise ValueError(
"zgrid bottom must match ground_altitude; "
f"expected {self.ground_altitude}, got {grid_bottom}"
)
grid_top = self.zgrid.levels[-1]
if not np.isclose(grid_top, self.toa_altitude):
raise ValueError(
"zgrid top must match toa_altitude; "
f"expected {self.toa_altitude}, got {grid_top}"
)
[docs]
@classmethod
def convert(cls, value: t.Any) -> t.Any:
"""
Attempt conversion of a value to a :class:`.SceneGeometry` subtype.
Parameters
----------
value
Value to attempt conversion of. If a dictionary is passed, its
``"type"`` key is used to route its other entries as keyword
arguments to the appropriate subtype's constructor. If a string is
passed, this method calls itself with the parameter
``{"type": value}``.
Returns
-------
result
If `value` is a dictionary, the constructed :class:`.SceneGeometry`
instance is returned. Otherwise, `value` is returned.
Raises
------
ValueError
A dictionary was passed but the requested type is unknown.
"""
if isinstance(value, str):
return cls.convert({"type": value})
if isinstance(value, dict):
value = value.copy()
geometry_type = value.pop("type")
# Note: if this conditional becomes large, use a dictionary
if geometry_type == "plane_parallel":
geometry_cls = PlaneParallelGeometry
elif geometry_type == "spherical_shell":
geometry_cls = SphericalShellGeometry
else:
raise ValueError(f"unknown geometry type '{geometry_type}'")
return geometry_cls(**pinttr.interpret_units(value, ureg=ureg))
return value
@property
@abstractmethod
def atmosphere_shape(self) -> Shape:
"""
:class:`.Shape`: Stencil of the participating medium representing the
atmosphere.
"""
pass
@property
@abstractmethod
def atmosphere_volume_to_world(self) -> mi.ScalarTransform4f:
"""
:class:`mi.ScalarTransform4f` : Mitsuba transform mapping volume texture
coordinates to world coordinates for heterogeneous atmosphere
components.
"""
pass
@property
@abstractmethod
def surface_shape(self) -> Shape:
"""
:class:`.Shape` : Shape representing the surface.
"""
pass
[docs]
@parse_docs
@attrs.define
class PlaneParallelGeometry(SceneGeometry):
"""
Plane parallel geometry.
A plane parallel atmosphere is translation-invariant in the X and Y
directions. However, Eradiate represents it with a finite 3D geometry
consisting of a cuboid. By default, the cuboid's size is computed
automatically; however, it can also be forced by assigning a value to
the `width` field.
"""
width: pint.Quantity = documented(
pinttr.field(default=1e6 * ureg.km, units=ucc.deferred("length")),
doc="Cuboid shape width.",
type="quantity",
init_type="quantity or float",
default="1,000,000 km",
)
@property
def atmosphere_shape(self) -> CuboidShape:
return CuboidShape.atmosphere(
top=self.toa_altitude, bottom=0.0 * ureg.km, width=self.width
)
@property
def atmosphere_volume_to_world(self) -> mi.ScalarTransform4f:
length_units = uck.get("length")
shape = self.atmosphere_shape
# The bounding box corresponds to the vertices of the cuboid
center = shape.center.m_as(length_units)
edges = shape.edges.m_as(length_units)
xmin, ymin = center[0:2] - edges[0:2] * 0.5
xmax, ymax = center[0:2] + edges[0:2] * 0.5
zmin = self.ground_altitude.m_as(length_units)
zmax = self.toa_altitude.m_as(length_units)
return map_unit_cube(xmin, xmax, ymin, ymax, zmin, zmax)
@property
def surface_shape(self) -> RectangleShape:
return RectangleShape.surface(altitude=self.ground_altitude, width=self.width)
[docs]
@parse_docs
@attrs.define
class SphericalShellGeometry(SceneGeometry):
"""
Spherical shell geometry.
A spherical shell atmosphere has a spherical symmetry. Eradiate represents
it with a finite 3D geometry consisting of a sphere. By default, the
sphere's radius is set equal to Earth's radius.
"""
planet_radius: pint.Quantity = documented(
pinttr.field(default=EARTH_RADIUS, units=ucc.deferred("length")),
doc="Planet radius. Defaults to Earth's radius.",
type="quantity",
init_type="quantity or float",
default=":data:`.EARTH_RADIUS`",
)
@property
def atmosphere_shape(self) -> SphereShape:
return SphereShape.atmosphere(
top=self.toa_altitude, planet_radius=self.planet_radius
)
@property
def atmosphere_volume_to_world(self) -> mi.ScalarTransform4f:
length_units = ucc.get("length")
# The bounding box corresponds to the vertices of the bounding box
bbox = self.atmosphere_shape.bbox
xmin, ymin, zmin = bbox.min.m_as(length_units)
xmax, ymax, zmax = bbox.max.m_as(length_units)
return map_cube(xmin, xmax, ymin, ymax, zmin, zmax)
@property
def atmosphere_volume_rmin(self) -> float:
"""
Value for the ``rmin`` parameter of the ``sphericalcoordsvolume`` plugin
describing the volume data in the spherical shell geometry.
"""
return (self.planet_radius / (self.planet_radius + self.toa_altitude)).m_as(
ureg.dimensionless
)
@property
def surface_shape(self) -> Shape:
return SphereShape.surface(
altitude=self.ground_altitude, planet_radius=self.planet_radius
)
