"""Declares useful sympy units and functions, to make it easier to work with `sympy` as the basis for a unit-aware system. Attributes: UNIT_BY_SYMBOL: Maps all abbreviated Sympy symbols to their corresponding Sympy unit. This is essential for parsing string expressions that use units, since a pure parse of ex. `a*m + m` would not otherwise be able to differentiate between `sp.Symbol(m)` and `spu.meter`. SympyType: A simple union of valid `sympy` types, used to check whether arbitrary objects should be handled using `sympy` functions. For simple `isinstance` checks, this should be preferred, as it is most performant. For general use, `SympyExpr` should be preferred. SympyExpr: A `SympyType` that is compatible with `pydantic`, including serialization/deserialization. Should be used via the `ConstrSympyExpr`, which also adds expression validation. """ import enum import functools import typing as typ from fractions import Fraction import jax import jax.numpy as jnp import pydantic as pyd import sympy as sp import sympy.physics.units as spu import typing_extensions as typx from pydantic_core import core_schema as pyd_core_schema from blender_maxwell import contracts as ct SympyType = ( sp.Basic | sp.Expr | sp.MatrixBase | sp.MutableDenseMatrix | spu.Quantity | spu.Dimension ) #################### # - Math Type #################### class MathType(enum.StrEnum): """Type identifiers that encompass common sets of mathematical objects.""" Bool = enum.auto() Integer = enum.auto() Rational = enum.auto() Real = enum.auto() Complex = enum.auto() def combine(*mathtypes: list[typ.Self]) -> typ.Self: if MathType.Complex in mathtypes: return MathType.Complex elif MathType.Real in mathtypes: return MathType.Real elif MathType.Rational in mathtypes: return MathType.Rational elif MathType.Integer in mathtypes: return MathType.Integer elif MathType.Bool in mathtypes: return MathType.Bool @staticmethod def from_expr(sp_obj: SympyType) -> type: ## TODO: Support for sp.Matrix if isinstance(sp_obj, sp.logic.boolalg.Boolean): return MathType.Bool if sp_obj.is_integer: return MathType.Integer if sp_obj.is_rational: return MathType.Rational if sp_obj.is_real: return MathType.Real if sp_obj.is_complex: return MathType.Complex msg = f"Can't determine MathType from sympy object: {sp_obj}" raise ValueError(msg) @staticmethod def from_pytype(dtype) -> type: return { bool: MathType.Bool, int: MathType.Integer, float: MathType.Real, complex: MathType.Complex, }[dtype] @property def pytype(self) -> type: MT = MathType return { MT.Bool: bool, MT.Integer: int, MT.Rational: float, MT.Real: float, MT.Complex: complex, }[self] @staticmethod def to_str(value: typ.Self) -> type: return { MathType.Bool: 'T|F', MathType.Integer: 'ℤ', MathType.Rational: 'ℚ', MathType.Real: 'ℝ', MathType.Complex: 'ℂ', }[value] @staticmethod def to_name(value: typ.Self) -> str: return MathType.to_str(value) @staticmethod def to_icon(value: typ.Self) -> str: return '' def bl_enum_element(self, i: int) -> ct.BLEnumElement: return ( str(self), MathType.to_name(self), MathType.to_name(self), MathType.to_icon(self), i, ) class NumberSize1D(enum.StrEnum): """Valid 1D-constrained shape.""" Scalar = enum.auto() Vec2 = enum.auto() Vec3 = enum.auto() Vec4 = enum.auto() @staticmethod def to_name(value: typ.Self) -> str: NS = NumberSize1D return { NS.Scalar: 'Scalar', NS.Vec2: '2D', NS.Vec3: '3D', NS.Vec4: '4D', }[value] @staticmethod def to_icon(value: typ.Self) -> str: NS = NumberSize1D return { NS.Scalar: '', NS.Vec2: '', NS.Vec3: '', NS.Vec4: '', }[value] def bl_enum_element(self, i: int) -> ct.BLEnumElement: return ( str(self), NumberSize1D.to_name(self), NumberSize1D.to_name(self), NumberSize1D.to_icon(self), i, ) @staticmethod def supports_shape(shape: tuple[int, ...] | None): return shape is None or (len(shape) == 1 and shape[0] in [2, 3]) @staticmethod def from_shape(shape: tuple[typ.Literal[2, 3]] | None) -> typ.Self: NS = NumberSize1D return { None: NS.Scalar, (2,): NS.Vec2, (3,): NS.Vec3, (4,): NS.Vec3, }[shape] @property def shape(self): NS = NumberSize1D return { NS.Scalar: None, NS.Vec2: (2,), NS.Vec3: (3,), NS.Vec3: (4,), }[self] #################### # - Unit Dimensions #################### class DimsMeta(type): def __getattr__(cls, attr: str) -> spu.Dimension: if ( attr in spu.definitions.dimension_definitions.__dir__() and not attr.startswith('__') ): return getattr(spu.definitions.dimension_definitions, attr) raise AttributeError(name=attr, obj=Dims) class Dims(metaclass=DimsMeta): """Access `sympy.physics.units` dimensions with less hassle. Any unit dimension available in `sympy.physics.units.definitions.dimension_definitions` can be accessed as an attribute of `Dims`. An `AttributeError` is raised if the unit cannot be found in `sympy`. Examples: The objects returned are a direct alias to `sympy`, with less hassle: ```python assert Dims.length == ( sympy.physics.units.definitions.dimension_definitions.length ) ``` """ #################### # - Units #################### femtosecond = fs = spu.Quantity('femtosecond', abbrev='fs') femtosecond.set_global_relative_scale_factor(spu.femto, spu.second) # Length femtometer = fm = spu.Quantity('femtometer', abbrev='fm') femtometer.set_global_relative_scale_factor(spu.femto, spu.meter) # Lum Flux lumen = lm = spu.Quantity('lumen', abbrev='lm') lumen.set_global_relative_scale_factor(1, spu.candela * spu.steradian) # Force nanonewton = nN = spu.Quantity('nanonewton', abbrev='nN') # noqa: N816 nanonewton.set_global_relative_scale_factor(spu.nano, spu.newton) micronewton = uN = spu.Quantity('micronewton', abbrev='μN') # noqa: N816 micronewton.set_global_relative_scale_factor(spu.micro, spu.newton) millinewton = mN = spu.Quantity('micronewton', abbrev='mN') # noqa: N816 micronewton.set_global_relative_scale_factor(spu.milli, spu.newton) # Frequency kilohertz = KHz = spu.Quantity('kilohertz', abbrev='KHz') kilohertz.set_global_relative_scale_factor(spu.kilo, spu.hertz) megahertz = MHz = spu.Quantity('megahertz', abbrev='MHz') kilohertz.set_global_relative_scale_factor(spu.kilo, spu.hertz) gigahertz = GHz = spu.Quantity('gigahertz', abbrev='GHz') gigahertz.set_global_relative_scale_factor(spu.giga, spu.hertz) terahertz = THz = spu.Quantity('terahertz', abbrev='THz') terahertz.set_global_relative_scale_factor(spu.tera, spu.hertz) petahertz = PHz = spu.Quantity('petahertz', abbrev='PHz') petahertz.set_global_relative_scale_factor(spu.peta, spu.hertz) exahertz = EHz = spu.Quantity('exahertz', abbrev='EHz') exahertz.set_global_relative_scale_factor(spu.exa, spu.hertz) # Pressure millibar = mbar = spu.Quantity('millibar', abbrev='mbar') millibar.set_global_relative_scale_factor(spu.milli, spu.bar) hectopascal = hPa = spu.Quantity('hectopascal', abbrev='hPa') # noqa: N816 hectopascal.set_global_relative_scale_factor(spu.hecto, spu.pascal) UNIT_BY_SYMBOL: dict[sp.Symbol, spu.Quantity] = { unit.name: unit for unit in spu.__dict__.values() if isinstance(unit, spu.Quantity) } | {unit.name: unit for unit in globals().values() if isinstance(unit, spu.Quantity)} #################### # - Expr Analysis: Units #################### ## TODO: Caching w/srepr'ed expression. ## TODO: An LFU cache could do better than an LRU. def uses_units(sp_obj: SympyType) -> bool: """Determines if an expression uses any units. Notes: The expression graph is traversed depth-first with `sp.postorder_traversal`, to search for `sp.Quantity` elements. Depth-first was chosen since `sp.Quantity`s are likelier to be found among individual symbols, rather than complete subexpressions. The **worst-case** runtime is when there are no units, in which case the **entire expression graph will be traversed**. Parameters: expr: The sympy expression that may contain units. Returns: Whether or not there are units used within the expression. """ return sp_obj.has(spu.Quantity) # return any( # isinstance(subexpr, spu.Quantity) for subexpr in sp.postorder_traversal(sp_obj) # ) ## TODO: Caching w/srepr'ed expression. ## TODO: An LFU cache could do better than an LRU. def get_units(expr: sp.Expr) -> set[spu.Quantity]: """Finds all units used by the expression, and returns them as a set. No information about _the relationship between units_ is exposed. For example, compound units like `spu.meter / spu.second` would be mapped to `{spu.meter, spu.second}`. Notes: The expression graph is traversed depth-first with `sp.postorder_traversal`, to search for `sp.Quantity` elements. The performance is comparable to the performance of `sp.postorder_traversal`, since the **entire expression graph will always be traversed**, with the added overhead of one `isinstance` call per expression-graph-node. Parameters: expr: The sympy expression that may contain units. Returns: All units (`spu.Quantity`) used within the expression. """ return { subexpr for subexpr in sp.postorder_traversal(expr) if isinstance(subexpr, spu.Quantity) } def parse_shape(sp_obj: SympyType) -> int | None: if isinstance(sp_obj, sp.MatrixBase): return sp_obj.shape return None #################### # - Pydantic-Validated SympyExpr #################### class _SympyExpr: """Low-level `pydantic`, schema describing how to serialize/deserialize fields that have a `SympyType` (like `sp.Expr`), so we can cleanly use `sympy` types in `pyd.BaseModel`. Notes: You probably want to use `SympyExpr`. Examples: To be usable as a type annotation on `pyd.BaseModel`, attach this to `SympyType` using `typx.Annotated`: ```python SympyExpr = typx.Annotated[SympyType, _SympyExpr] class Spam(pyd.BaseModel): line: SympyExpr = sp.Eq(sp.y, 2*sp.Symbol(x, real=True) - 3) ``` """ @classmethod def __get_pydantic_core_schema__( cls, _source_type: SympyType, _handler: pyd.GetCoreSchemaHandler, ) -> pyd_core_schema.CoreSchema: """Compute a schema that allows `pydantic` to validate a `sympy` type.""" def validate_from_str(sp_str: str | typ.Any) -> SympyType | typ.Any: """Parse and validate a string expression. Parameters: sp_str: A stringified `sympy` object, that will be parsed to a sympy type. Before use, `isinstance(expr_str, str)` is checked. If the object isn't a string, then the validation will be skipped. Returns: Either a `sympy` object, if the input is parseable, or the same untouched object. Raises: ValueError: If `sp_str` is a string, but can't be parsed into a `sympy` expression. """ # Constrain to String if not isinstance(sp_str, str): return sp_str # Parse String -> Sympy try: expr = sp.sympify(sp_str) except ValueError as ex: msg = f'String {sp_str} is not a valid sympy expression' raise ValueError(msg) from ex # Substitute Symbol -> Quantity return expr.subs(UNIT_BY_SYMBOL) def validate_from_pytype( sp_pytype: int | Fraction | float | complex, ) -> SympyType | typ.Any: """Parse and validate a pure Python type. Parameters: sp_str: A stringified `sympy` object, that will be parsed to a sympy type. Before use, `isinstance(expr_str, str)` is checked. If the object isn't a string, then the validation will be skipped. Returns: Either a `sympy` object, if the input is parseable, or the same untouched object. Raises: ValueError: If `sp_str` is a string, but can't be parsed into a `sympy` expression. """ # Constrain to String if not isinstance(sp_pytype, int | Fraction | float | complex): return sp_pytype if isinstance(sp_pytype, int): return sp.Integer(sp_pytype) if isinstance(sp_pytype, Fraction): return sp.Rational(sp_pytype.numerator, sp_pytype.denominator) if isinstance(sp_pytype, float): return sp.Float(sp_pytype) # sp_pytype => Complex return sp_pytype.real + sp.I * sp_pytype.imag sympy_expr_schema = pyd_core_schema.chain_schema( [ pyd_core_schema.no_info_plain_validator_function(validate_from_str), pyd_core_schema.no_info_plain_validator_function(validate_from_pytype), pyd_core_schema.is_instance_schema(SympyType), ] ) return pyd_core_schema.json_or_python_schema( json_schema=sympy_expr_schema, python_schema=sympy_expr_schema, serialization=pyd_core_schema.plain_serializer_function_ser_schema( lambda sp_obj: sp.srepr(sp_obj) ), ) SympyExpr = typx.Annotated[ sp.Basic, ## Treat all sympy types as sp.Basic _SympyExpr, ] ## TODO: The type game between SympyType, SympyExpr, and the various flavors of ConstrSympyExpr(), is starting to be a bit much. Let's consolidate. def ConstrSympyExpr( # noqa: N802, PLR0913 # Features allow_variables: bool = True, allow_units: bool = True, # Structures allowed_sets: set[typ.Literal['integer', 'rational', 'real', 'complex']] | None = None, allowed_structures: set[typ.Literal['scalar', 'matrix']] | None = None, # Element Class max_symbols: int | None = None, allowed_symbols: set[sp.Symbol] | None = None, allowed_units: set[spu.Quantity] | None = None, # Shape Class allowed_matrix_shapes: set[tuple[int, int]] | None = None, ) -> SympyType: """Constructs a `SympyExpr` type, which will validate `sympy` types when used in a `pyd.BaseModel`. Relies on the `sympy` assumptions system. See Parameters (TBD): Returns: A type that represents a constrained `sympy` expression. """ def validate_expr(expr: SympyType): if not (isinstance(expr, SympyType),): msg = f"expr '{expr}' is not an allowed Sympy expression ({SympyType})" raise ValueError(msg) msgs = set() # Validate Feature Class if (not allow_variables) and (len(expr.free_symbols) > 0): msgs.add( f'allow_variables={allow_variables} does not match expression {expr}.' ) if (not allow_units) and uses_units(expr): msgs.add(f'allow_units={allow_units} does not match expression {expr}.') # Validate Structure Class if ( allowed_sets and isinstance(expr, sp.Expr) and not any( { 'integer': expr.is_integer, 'rational': expr.is_rational, 'real': expr.is_real, 'complex': expr.is_complex, }[allowed_set] for allowed_set in allowed_sets ) ): msgs.add( f"allowed_sets={allowed_sets} does not match expression {expr} (remember to add assumptions to symbols, ex. `x = sp.Symbol('x', real=True))" ) if allowed_structures and not any( { 'scalar': True, 'matrix': isinstance(expr, sp.MatrixBase), }[allowed_set] for allowed_set in allowed_structures ): msgs.add( f"allowed_structures={allowed_structures} does not match expression {expr} (remember to add assumptions to symbols, ex. `x = sp.Symbol('x', real=True))" ) # Validate Element Class if max_symbols and len(expr.free_symbols) > max_symbols: msgs.add(f'max_symbols={max_symbols} does not match expression {expr}') if allowed_symbols and expr.free_symbols.issubset(allowed_symbols): msgs.add( f'allowed_symbols={allowed_symbols} does not match expression {expr}' ) if allowed_units and get_units(expr).issubset(allowed_units): msgs.add(f'allowed_units={allowed_units} does not match expression {expr}') # Validate Shape Class if ( allowed_matrix_shapes and isinstance(expr, sp.MatrixBase) ) and expr.shape not in allowed_matrix_shapes: msgs.add( f'allowed_matrix_shapes={allowed_matrix_shapes} does not match expression {expr} with shape {expr.shape}' ) # Error or Return if msgs: raise ValueError(str(msgs)) return expr return typx.Annotated[ sp.Basic, _SympyExpr, pyd.AfterValidator(validate_expr), ] #################### # - Common ConstrSympyExpr #################### # Expression ScalarUnitlessRealExpr: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=False, allowed_structures={'scalar'}, allowed_sets={'integer', 'rational', 'real'}, ) ScalarUnitlessComplexExpr: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=False, allowed_structures={'scalar'}, allowed_sets={'integer', 'rational', 'real', 'complex'}, ) # Symbol IntSymbol: typ.TypeAlias = ConstrSympyExpr( allow_variables=True, allow_units=False, allowed_sets={'integer'}, max_symbols=1, ) RationalSymbol: typ.TypeAlias = ConstrSympyExpr( allow_variables=True, allow_units=False, allowed_sets={'integer', 'rational'}, max_symbols=1, ) RealSymbol: typ.TypeAlias = ConstrSympyExpr( allow_variables=True, allow_units=False, allowed_sets={'integer', 'rational', 'real'}, max_symbols=1, ) ComplexSymbol: typ.TypeAlias = ConstrSympyExpr( allow_variables=True, allow_units=False, allowed_sets={'integer', 'rational', 'real', 'complex'}, max_symbols=1, ) Symbol: typ.TypeAlias = IntSymbol | RealSymbol | ComplexSymbol # Unit UnitDimension: typ.TypeAlias = SympyExpr ## Actually spu.Dimension ## Technically a "unit expression", which includes compound types. ## Support for this is the reason to prefer over raw spu.Quantity. Unit: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=True, allowed_structures={'scalar'}, ) # Number IntNumber: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=False, allowed_sets={'integer'}, allowed_structures={'scalar'}, ) RealNumber: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=False, allowed_sets={'integer', 'rational', 'real'}, allowed_structures={'scalar'}, ) ComplexNumber: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=False, allowed_sets={'integer', 'rational', 'real', 'complex'}, allowed_structures={'scalar'}, ) Number: typ.TypeAlias = IntNumber | RealNumber | ComplexNumber # Number PhysicalRealNumber: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=True, allowed_sets={'integer', 'rational', 'real'}, allowed_structures={'scalar'}, ) PhysicalComplexNumber: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=True, allowed_sets={'integer', 'rational', 'real', 'complex'}, allowed_structures={'scalar'}, ) PhysicalNumber: typ.TypeAlias = PhysicalRealNumber | PhysicalComplexNumber # Vector Real3DVector: typ.TypeAlias = ConstrSympyExpr( allow_variables=False, allow_units=False, allowed_sets={'integer', 'rational', 'real'}, allowed_structures={'matrix'}, allowed_matrix_shapes={(3, 1)}, ) #################### # - Sympy Utilities: Printing #################### _SYMPY_EXPR_PRINTER_STR = sp.printing.str.StrPrinter( settings={ 'abbrev': True, } ) def sp_to_str(sp_obj: SympyExpr) -> str: """Converts a sympy object to an output-oriented string (w/abbreviated units), using a dedicated StrPrinter. This should be used whenever a **string for UI use** is needed from a `sympy` object. Notes: This should **NOT** be used in cases where the string will be `sp.sympify()`ed back into a sympy expression. For such cases, rely on `sp.srepr()`, which uses an _explicit_ representation. Parameters: sp_obj: The `sympy` object to convert to a string. Returns: A string representing the expression for human use. _The string is not re-encodable to the expression._ """ return _SYMPY_EXPR_PRINTER_STR.doprint(sp_obj) def pretty_symbol(sym: sp.Symbol) -> str: return f'{sym.name} ∈ ' + ( 'ℂ' if sym.is_complex else ('ℝ' if sym.is_real else ('ℤ' if sym.is_integer else '?')) ) #################### # - Unit Utilities #################### def scale_to_unit(sp_obj: SympyType, unit: spu.Quantity) -> Number: """Convert an expression that uses units to a different unit, then strip all units, leaving only a unitless `sympy` value. This is used whenever the unitless part of an expression is needed, but guaranteed expressed in a particular unit, aka. **unit system normalization**. Notes: The unitless output is still an `sp.Expr`, which may contain ex. symbols. If you know that the output **should** work as a corresponding Python type (ex. `sp.Integer` vs. `int`), but it doesn't, you can use `sympy_to_python()` to produce a pure-Python type. In this way, with a little care, broad compatiblity can be bridged between the `sympy.physics.units` unit system and the wider Python ecosystem. Parameters: expr: The unit-containing expression to convert. unit_to: The unit that is converted to. Returns: The unitless part of `expr`, after scaling the entire expression to `unit`. Raises: ValueError: If the result of unit-conversion and -stripping still has units, as determined by `uses_units()`. """ ## TODO: An LFU cache could do better than an LRU. unitless_expr = spu.convert_to(sp_obj, unit) / unit if not uses_units(unitless_expr): return unitless_expr msg = f'Sympy object "{sp_obj}" was scaled to the unit "{unit}" with the expectation that the result would be unitless, but the result "{unitless_expr}" has units "{get_units(unitless_expr)}"' raise ValueError(msg) def scaling_factor(unit_from: spu.Quantity, unit_to: spu.Quantity) -> Number: """Compute the numerical scaling factor imposed on the unitless part of the expression when converting from one unit to another. Parameters: unit_from: The unit that is converted from. unit_to: The unit that is converted to. Returns: The numerical scaling factor between the two units. Raises: ValueError: If the two units don't share a common dimension. """ if unit_from.dimension == unit_to.dimension: return scale_to_unit(unit_from, unit_to) msg = f"Dimension of unit_from={unit_from} ({unit_from.dimension}) doesn't match the dimension of unit_to={unit_to} ({unit_to.dimension}); therefore, there is no scaling factor between them" raise ValueError(msg) _UNIT_STR_MAP = {sym.name: unit for sym, unit in UNIT_BY_SYMBOL.items()} @functools.cache def unit_str_to_unit(unit_str: str) -> Unit | None: if unit_str in _UNIT_STR_MAP: return _UNIT_STR_MAP[unit_str] msg = 'No valid unit for unit string {unit_str}' raise ValueError(msg) #################### # - "Physical" Type #################### class PhysicalType(enum.StrEnum): """Type identifiers for expressions with both `MathType` and a unit, aka a "physical" type.""" # Global Time = enum.auto() Angle = enum.auto() SolidAngle = enum.auto() ## TODO: Some kind of 3D-specific orientation ex. a quaternion Freq = enum.auto() AngFreq = enum.auto() ## rad*hertz # Cartesian Length = enum.auto() Area = enum.auto() Volume = enum.auto() # Mechanical Vel = enum.auto() Accel = enum.auto() Mass = enum.auto() Force = enum.auto() Pressure = enum.auto() # Energy Work = enum.auto() ## joule Power = enum.auto() ## watt PowerFlux = enum.auto() ## watt Temp = enum.auto() # Electrodynamics Current = enum.auto() ## ampere CurrentDensity = enum.auto() Charge = enum.auto() ## coulomb Voltage = enum.auto() Capacitance = enum.auto() ## farad Impedance = enum.auto() ## ohm Conductance = enum.auto() ## siemens Conductivity = enum.auto() ## siemens / length MFlux = enum.auto() ## weber MFluxDensity = enum.auto() ## tesla Inductance = enum.auto() ## henry EField = enum.auto() HField = enum.auto() # Luminal LumIntensity = enum.auto() LumFlux = enum.auto() Luminance = enum.auto() Illuminance = enum.auto() # Optics OrdinaryWaveVector = enum.auto() AngularWaveVector = enum.auto() PoyntingVector = enum.auto() @property def unit_dim(self): PT = PhysicalType return { # Global PT.Time: Dims.time, PT.Angle: Dims.angle, PT.SolidAngle: Dims.steradian, ## MISSING PT.Freq: Dims.frequency, PT.AngFreq: Dims.angle * Dims.frequency, # Cartesian PT.Length: Dims.length, PT.Area: Dims.length**2, PT.Volume: Dims.length**3, # Mechanical PT.Vel: Dims.length / Dims.time, PT.Accel: Dims.length / Dims.time**2, PT.Mass: Dims.mass, PT.Force: Dims.force, PT.Pressure: Dims.pressure, # Energy PT.Work: Dims.energy, PT.Power: Dims.power, PT.PowerFlux: Dims.power / Dims.length**2, PT.Temp: Dims.temperature, # Electrodynamics PT.Current: Dims.current, PT.CurrentDensity: Dims.current / Dims.length**2, PT.Charge: Dims.charge, PT.Voltage: Dims.voltage, PT.Capacitance: Dims.capacitance, PT.Impedance: Dims.impedance, PT.Conductance: Dims.conductance, PT.Conductivity: Dims.conductance / Dims.length, PT.MFlux: Dims.magnetic_flux, PT.MFluxDensity: Dims.magnetic_density, PT.Inductance: Dims.inductance, PT.EField: Dims.voltage / Dims.length, PT.HField: Dims.current / Dims.length, # Luminal PT.LumIntensity: Dims.luminous_intensity, PT.LumFlux: Dims.luminous_intensity * Dims.steradian, PT.Illuminance: Dims.luminous_intensity / Dims.length**2, # Optics PT.OrdinaryWaveVector: Dims.frequency, PT.AngularWaveVector: Dims.angle * Dims.frequency, PT.PoyntingVector: Dims.power / Dims.length**2, } @property def default_unit(self) -> list[Unit]: PT = PhysicalType return { # Global PT.Time: spu.picosecond, PT.Angle: spu.radian, PT.SolidAngle: spu.steradian, PT.Freq: terahertz, PT.AngFreq: spu.radian * terahertz, # Cartesian PT.Length: spu.micrometer, PT.Area: spu.um**2, PT.Volume: spu.um**3, # Mechanical PT.Vel: spu.um / spu.second, PT.Accel: spu.um / spu.second, PT.Mass: spu.microgram, PT.Force: micronewton, PT.Pressure: millibar, # Energy PT.Work: spu.joule, PT.Power: spu.watt, PT.PowerFlux: spu.watt / spu.meter**2, PT.Temp: spu.kelvin, # Electrodynamics PT.Current: spu.ampere, PT.CurrentDensity: spu.ampere / spu.meter**2, PT.Charge: spu.coulomb, PT.Voltage: spu.volt, PT.Capacitance: spu.farad, PT.Impedance: spu.ohm, PT.Conductance: spu.siemens, PT.Conductivity: spu.siemens / spu.micrometer, PT.MFlux: spu.weber, PT.MFluxDensity: spu.tesla, PT.Inductance: spu.henry, PT.EField: spu.volt / spu.micrometer, PT.HField: spu.ampere / spu.micrometer, # Luminal PT.LumIntensity: spu.candela, PT.LumFlux: spu.candela * spu.steradian, PT.Illuminance: spu.candela / spu.meter**2, # Optics PT.OrdinaryWaveVector: terahertz, PT.AngularWaveVector: spu.radian * terahertz, }[self] @property def valid_units(self) -> list[Unit]: PT = PhysicalType return { # Global PT.Time: [ femtosecond, spu.picosecond, spu.nanosecond, spu.microsecond, spu.millisecond, spu.second, spu.minute, spu.hour, spu.day, ], PT.Angle: [ spu.radian, spu.degree, ], PT.SolidAngle: [ spu.steradian, ], PT.Freq: ( _valid_freqs := [ spu.hertz, kilohertz, megahertz, gigahertz, terahertz, petahertz, exahertz, ] ), PT.AngFreq: [spu.radian * _unit for _unit in _valid_freqs], # Cartesian PT.Length: ( _valid_lens := [ spu.picometer, spu.angstrom, spu.nanometer, spu.micrometer, spu.millimeter, spu.centimeter, spu.meter, spu.inch, spu.foot, spu.yard, spu.mile, ] ), PT.Area: [_unit**2 for _unit in _valid_lens], PT.Volume: [_unit**3 for _unit in _valid_lens], # Mechanical PT.Vel: [_unit / spu.second for _unit in _valid_lens], PT.Accel: [_unit / spu.second**2 for _unit in _valid_lens], PT.Mass: [ spu.electron_rest_mass, spu.dalton, spu.microgram, spu.milligram, spu.gram, spu.kilogram, spu.metric_ton, ], PT.Force: [ spu.kg * spu.meter / spu.second**2, nanonewton, micronewton, millinewton, spu.newton, ], PT.Pressure: [ millibar, spu.bar, spu.pascal, hectopascal, spu.atmosphere, spu.psi, spu.mmHg, spu.torr, ], # Energy PT.Work: [ spu.electronvolt, spu.joule, ], PT.Power: [ spu.watt, ], PT.PowerFlux: [ spu.watt / spu.meter**2, ], PT.Temp: [ spu.kelvin, ], # Electrodynamics PT.Current: [ spu.ampere, ], PT.CurrentDensity: [ spu.ampere / spu.meter**2, ], PT.Charge: [ spu.coulomb, ], PT.Voltage: [ spu.volt, ], PT.Capacitance: [ spu.farad, ], PT.Impedance: [ spu.ohm, ], PT.Conductance: [ spu.siemens, ], PT.Conductivity: [ spu.siemens / spu.micrometer, spu.siemens / spu.meter, ], PT.MFlux: [ spu.weber, ], PT.MFluxDensity: [ spu.tesla, ], PT.Inductance: [ spu.henry, ], PT.EField: [ spu.volt / spu.micrometer, spu.volt / spu.meter, ], PT.HField: [ spu.ampere / spu.micrometer, spu.ampere / spu.meter, ], # Luminal PT.LumIntensity: [ spu.candela, ], PT.LumFlux: [ spu.candela * spu.steradian, ], PT.Illuminance: [ spu.candela / spu.meter**2, ], # Optics PT.OrdinaryWaveVector: _valid_freqs, PT.AngularWaveVector: [spu.radian * _unit for _unit in _valid_freqs], }[self] @staticmethod def from_unit(unit: Unit) -> list[Unit]: for physical_type in list[PhysicalType]: if unit in physical_type.valid_units: return physical_type msg = f'No PhysicalType found for unit {unit}' raise ValueError(msg) @property def valid_shapes(self): PT = PhysicalType overrides = { # Cartesian PT.Length: [None, (2,), (3,)], # Mechanical PT.Vel: [None, (2,), (3,)], PT.Accel: [None, (2,), (3,)], PT.Force: [None, (2,), (3,)], # Energy PT.Work: [None, (2,), (3,)], PT.PowerFlux: [None, (2,), (3,)], # Electrodynamics PT.CurrentDensity: [None, (2,), (3,)], PT.MFluxDensity: [None, (2,), (3,)], PT.EField: [None, (2,), (3,)], PT.HField: [None, (2,), (3,)], # Luminal PT.LumFlux: [None, (2,), (3,)], # Optics PT.OrdinaryWaveVector: [None, (2,), (3,)], PT.AngularWaveVector: [None, (2,), (3,)], PT.PoyntingVector: [None, (2,), (3,)], } return overrides.get(self, [None]) @property def valid_mathtypes(self) -> list[MathType]: """Returns a list of valid mathematical types, especially whether it can be real- or complex-valued. Generally, all unit quantities are real, in the algebraic mathematical sense. However, in electrodynamics especially, it becomes enormously useful to bake in a _rotational component_ as an imaginary value, be it simply to model phase or oscillation-oriented dampening. This imaginary part has physical meaning, which can be expressed using the same mathematical formalism associated with unit systems. In general, the value is a phasor. While it is difficult to arrive at a well-defined way of saying, "this is when a quantity is complex", an attempt has been made to form a sensible baseline based on when phasor math may apply. Notes: - **Freq**/**AngFreq**: The imaginary part represents growth/dampening of the oscillation. - **Current**/**Voltage**: The imaginary part represents the phase. This also holds for any downstream units. - **Charge**: Generally, it is real. However, an imaginary phase term seems to have research applications when dealing with high-order harmonics in high-energy pulsed lasers: - **Conductance**: The imaginary part represents the extinction, in the Drude-model sense. - **Poynting**: The imaginary part represents the oscillation in the power flux over time. """ MT = MathType PT = PhysicalType overrides = { # Cartesian PT.Freq: [MT.Real, MT.Complex], ## Im -> Growth/Damping PT.AngFreq: [MT.Real, MT.Complex], ## Im -> Growth/Damping # Mechanical # Energy # Electrodynamics PT.Current: [MT.Real, MT.Complex], ## Im -> Phase PT.CurrentDensity: [MT.Real, MT.Complex], ## Im -> Phase PT.Charge: [MT.Real, MT.Complex], ## Im -> Phase PT.Voltage: [MT.Real, MT.Complex], ## Im -> Phase PT.Capacitance: [MT.Real, MT.Complex], ## Im -> Phase PT.Impedance: [MT.Real, MT.Complex], ## Im -> Reactance PT.Inductance: [MT.Real, MT.Complex], ## Im -> Extinction PT.Conductance: [MT.Real, MT.Complex], ## Im -> Extinction PT.Conductivity: [MT.Real, MT.Complex], ## Im -> Extinction PT.MFlux: [MT.Real, MT.Complex], ## Im -> Phase PT.MFluxDensity: [MT.Real, MT.Complex], ## Im -> Phase PT.EField: [MT.Real, MT.Complex], ## Im -> Phase PT.HField: [MT.Real, MT.Complex], ## Im -> Phase # Luminal # Optics PT.OrdinaryWaveVector: [MT.Real, MT.Complex], ## Im -> Phase PT.AngularWaveVector: [MT.Real, MT.Complex], ## Im -> Phase PT.PoyntingVector: [MT.Real, MT.Complex], ## Im -> Reactive Power } return overrides.get(self, [MT.Real]) @staticmethod def to_name(value: typ.Self) -> str: return sp_to_str(value.unit_dim) @staticmethod def to_icon(value: typ.Self) -> str: return '' def bl_enum_element(self, i: int) -> ct.BLEnumElement: PT = PhysicalType return ( str(self), PT.to_name(self), PT.to_name(self), PT.to_icon(self), i, ) #################### # - Standard Unit Systems #################### UnitSystem: typ.TypeAlias = dict[PhysicalType, Unit] _PT = PhysicalType UNITS_SI: UnitSystem = { # Global _PT.Time: spu.second, _PT.Angle: spu.radian, _PT.SolidAngle: spu.steradian, _PT.Freq: spu.hertz, _PT.AngFreq: spu.radian * spu.hertz, # Cartesian _PT.Length: spu.meter, _PT.Area: spu.meter**2, _PT.Volume: spu.meter**3, # Mechanical _PT.Vel: spu.meter / spu.second, _PT.Accel: spu.meter / spu.second**2, _PT.Mass: spu.kilogram, _PT.Force: spu.newton, # Energy _PT.Work: spu.joule, _PT.Power: spu.watt, _PT.PowerFlux: spu.watt / spu.meter**2, _PT.Temp: spu.kelvin, # Electrodynamics _PT.Current: spu.ampere, _PT.CurrentDensity: spu.ampere / spu.meter**2, _PT.Capacitance: spu.farad, _PT.Impedance: spu.ohm, _PT.Conductance: spu.siemens, _PT.Conductivity: spu.siemens / spu.meter, _PT.MFlux: spu.weber, _PT.MFluxDensity: spu.tesla, _PT.Inductance: spu.henry, _PT.EField: spu.volt / spu.meter, _PT.HField: spu.ampere / spu.meter, # Luminal _PT.LumIntensity: spu.candela, _PT.LumFlux: lumen, _PT.Illuminance: spu.lux, # Optics _PT.OrdinaryWaveVector: spu.hertz, _PT.AngularWaveVector: spu.radian * spu.hertz, _PT.PoyntingVector: spu.watt / spu.meter**2, } #################### # - Sympy Utilities: Cast to Python #################### def sympy_to_python( scalar: sp.Basic, use_jax_array: bool = False ) -> int | float | complex | tuple | jax.Array: """Convert a scalar sympy expression to the directly corresponding Python type. Arguments: scalar: A sympy expression that has no symbols, but is expressed as a Sympy type. For expressions that are equivalent to a scalar (ex. "(2a + a)/a"), you must simplify the expression with ex. `sp.simplify()` before passing to this parameter. Returns: A pure Python type that directly corresponds to the input scalar expression. """ if isinstance(scalar, sp.MatrixBase): # Detect Single Column Vector ## --> Flatten to Single Row Vector if len(scalar.shape) == 2 and scalar.shape[1] == 1: _scalar = scalar.T else: _scalar = scalar # Convert to Tuple of Tuples matrix = tuple( [tuple([sympy_to_python(el) for el in row]) for row in _scalar.tolist()] ) # Detect Single Row Vector ## --> This could be because the scalar had it. ## --> This could also be because we flattened a column vector. ## Either way, we should strip the pointless dimensions. if len(matrix) == 1: return matrix[0] if not use_jax_array else jnp.array(matrix[0]) return matrix if not use_jax_array else jnp.array(matrix) if scalar.is_integer: return int(scalar) if scalar.is_rational or scalar.is_real: return float(scalar) if scalar.is_complex: return complex(scalar) msg = f'Cannot convert sympy scalar expression "{scalar}" to a Python type. Check the assumptions on the expr (current expr assumptions: "{scalar._assumptions}")' # noqa: SLF001 raise ValueError(msg) #################### # - Convert to Unit System #################### def _flat_unit_system_units(unit_system: UnitSystem) -> SympyExpr: return list(unit_system.values()) def convert_to_unit_system(sp_obj: SympyExpr, unit_system: UnitSystem) -> SympyExpr: """Convert an expression to the units of a given unit system, with appropriate scaling.""" return spu.convert_to(sp_obj, _flat_unit_system_units(unit_system)) def scale_to_unit_system( sp_obj: SympyExpr, unit_system: UnitSystem, use_jax_array: bool = False ) -> int | float | complex | tuple | jax.Array: """Convert an expression to the units of a given unit system, then strip all units of the unit system. Unit stripping is "dumb": Substitute any `sympy` object in `unit_system.values()` with `1`. Afterwards, it is converted to an appropriate Python type. Notes: For stability, and performance, reasons, this should only be used at the very last stage. Regarding performance: **This is not a fast function**. Parameters: sp_obj: An arbitrary sympy object, presumably with units. unit_system: A unit system mapping `PhysicalType` to particular choices of (compound) units. Note that, in this context, only `unit_system.values()` is used. Returns: An appropriate pure Python type, after scaling to the unit system and stripping all units away. If the returned type is array-like, and `use_jax_array` is specified, then (and **only** then) will a `jax.Array` be returned instead of a nested `tuple`. """ return sympy_to_python( convert_to_unit_system(sp_obj, unit_system).subs( {unit: 1 for unit in unit_system.values()} ), use_jax_array=use_jax_array, )