1696 lines
50 KiB
Python
1696 lines
50 KiB
Python
# blender_maxwell
|
||
# Copyright (C) 2024 blender_maxwell Project Contributors
|
||
#
|
||
# This program is free software: you can redistribute it and/or modify
|
||
# it under the terms of the GNU Affero General Public License as published by
|
||
# the Free Software Foundation, either version 3 of the License, or
|
||
# (at your option) any later version.
|
||
#
|
||
# This program is distributed in the hope that it will be useful,
|
||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||
# GNU Affero General Public License for more details.
|
||
#
|
||
# You should have received a copy of the GNU Affero General Public License
|
||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||
|
||
"""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 sys
|
||
import typing as typ
|
||
from fractions import Fraction
|
||
|
||
import jax
|
||
import jax.numpy as jnp
|
||
import jaxtyping as jtyp
|
||
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
|
||
|
||
from . import logger
|
||
from .staticproperty import staticproperty
|
||
|
||
log = logger.get(__name__)
|
||
|
||
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."""
|
||
|
||
Integer = enum.auto()
|
||
Rational = enum.auto()
|
||
Real = enum.auto()
|
||
Complex = enum.auto()
|
||
|
||
@staticmethod
|
||
def combine(*mathtypes: list[typ.Self], optional: bool = False) -> typ.Self | None:
|
||
if MathType.Complex in mathtypes:
|
||
return MathType.Complex
|
||
if MathType.Real in mathtypes:
|
||
return MathType.Real
|
||
if MathType.Rational in mathtypes:
|
||
return MathType.Rational
|
||
if MathType.Integer in mathtypes:
|
||
return MathType.Integer
|
||
|
||
if optional:
|
||
return None
|
||
|
||
msg = f"Can't combine mathtypes {mathtypes}"
|
||
raise ValueError(msg)
|
||
|
||
def is_compatible(self, other: typ.Self) -> bool:
|
||
MT = MathType
|
||
return (
|
||
other
|
||
in {
|
||
MT.Integer: [MT.Integer],
|
||
MT.Rational: [MT.Integer, MT.Rational],
|
||
MT.Real: [MT.Integer, MT.Rational, MT.Real],
|
||
MT.Complex: [MT.Integer, MT.Rational, MT.Real, MT.Complex],
|
||
}[self]
|
||
)
|
||
|
||
def coerce_compatible_pyobj(
|
||
self, pyobj: bool | int | Fraction | float | complex
|
||
) -> int | Fraction | float | complex:
|
||
MT = MathType
|
||
match self:
|
||
case MT.Integer:
|
||
return int(pyobj)
|
||
case MT.Rational if isinstance(pyobj, int):
|
||
return Fraction(pyobj, 1)
|
||
case MT.Rational if isinstance(pyobj, Fraction):
|
||
return pyobj
|
||
case MT.Real:
|
||
return float(pyobj)
|
||
case MT.Complex if isinstance(pyobj, int | Fraction):
|
||
return complex(float(pyobj), 0)
|
||
case MT.Complex if isinstance(pyobj, float):
|
||
return complex(pyobj, 0)
|
||
|
||
@staticmethod
|
||
def from_expr(sp_obj: SympyType, optional: bool = False) -> type | None:
|
||
if isinstance(sp_obj, sp.MatrixBase):
|
||
return MathType.combine(
|
||
*[MathType.from_expr(v) for v in sp.flatten(sp_obj)]
|
||
)
|
||
|
||
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
|
||
|
||
# Infinities
|
||
if sp_obj in [sp.oo, -sp.oo]:
|
||
return MathType.Real ## TODO: Strictly, could be ex. integer...
|
||
if sp_obj in [sp.zoo, -sp.zoo]:
|
||
return MathType.Complex
|
||
|
||
if optional:
|
||
return None
|
||
|
||
msg = f"Can't determine MathType from sympy object: {sp_obj}"
|
||
raise ValueError(msg)
|
||
|
||
@staticmethod
|
||
def from_pytype(dtype: type) -> type:
|
||
return {
|
||
int: MathType.Integer,
|
||
Fraction: MathType.Rational,
|
||
float: MathType.Real,
|
||
complex: MathType.Complex,
|
||
}[dtype]
|
||
|
||
@staticmethod
|
||
def from_jax_array(data: jtyp.Shaped[jtyp.Array, '...']) -> type:
|
||
"""Deduce the MathType corresponding to a JAX array.
|
||
|
||
We go about this by leveraging that:
|
||
- `data` is of a homogeneous type.
|
||
- `data.item(0)` returns a single element of the array w/pure-python type.
|
||
|
||
By combing this with `type()` and `MathType.from_pytype`, we can effectively deduce the `MathType` of the entire array with relative efficiency.
|
||
|
||
Notes:
|
||
Should also work with numpy arrays.
|
||
"""
|
||
return MathType.from_pytype(type(data.item(0)))
|
||
|
||
@staticmethod
|
||
def has_mathtype(obj: typ.Any) -> typ.Literal['pytype', 'expr'] | None:
|
||
if isinstance(obj, bool | int | Fraction | float | complex):
|
||
return 'pytype'
|
||
if isinstance(obj, sp.Basic | sp.MatrixBase | sp.MutableDenseMatrix):
|
||
return 'expr'
|
||
|
||
return None
|
||
|
||
@property
|
||
def pytype(self) -> type:
|
||
MT = MathType
|
||
return {
|
||
MT.Integer: int,
|
||
MT.Rational: Fraction,
|
||
MT.Real: float,
|
||
MT.Complex: complex,
|
||
}[self]
|
||
|
||
@property
|
||
def symbolic_set(self) -> type:
|
||
MT = MathType
|
||
return {
|
||
MT.Integer: sp.Integers,
|
||
MT.Rational: sp.Rationals,
|
||
MT.Real: sp.Reals,
|
||
MT.Complex: sp.Complexes,
|
||
}[self]
|
||
|
||
@property
|
||
def inf_finite(self) -> type:
|
||
"""Opinionated finite representation of "infinity" within this `MathType`.
|
||
|
||
These are chosen using `sys.maxsize` and `sys.float_info`.
|
||
As such, while not arbitrary, this "finite representation of infinity" certainly is opinionated.
|
||
|
||
**Note** that, in practice, most systems will have no trouble working with values that exceed those defined here.
|
||
|
||
Notes:
|
||
Values should be presumed to vary by-platform, as the `sys` attributes may be influenced by CPU architecture, OS, runtime environment, etc. .
|
||
|
||
These values can be used directly in `jax` arrays, but at the cost of an overflow warning (in part because `jax` generally only allows the use of `float32`).
|
||
In this case, the warning doesn't matter, as the value will be cast to `jnp.inf` anyway.
|
||
|
||
However, it's generally cleaner to directly use `jnp.inf` if infinite values must be defined in an array context.
|
||
"""
|
||
MT = MathType
|
||
Z = MT.Integer
|
||
R = MT.Integer
|
||
return {
|
||
MT.Integer: (-sys.maxsize, sys.maxsize),
|
||
MT.Rational: (
|
||
Fraction(Z.inf_finite[0], 1),
|
||
Fraction(Z.inf_finite[1], 1),
|
||
),
|
||
MT.Real: -(sys.float_info.min, sys.float_info.max),
|
||
MT.Complex: (
|
||
complex(R.inf_finite[0], R.inf_finite[0]),
|
||
complex(R.inf_finite[1], R.inf_finite[1]),
|
||
),
|
||
}[self]
|
||
|
||
@property
|
||
def sp_symbol_a(self) -> type:
|
||
MT = MathType
|
||
return {
|
||
MT.Integer: sp.Symbol('a', integer=True),
|
||
MT.Rational: sp.Symbol('a', rational=True),
|
||
MT.Real: sp.Symbol('a', real=True),
|
||
MT.Complex: sp.Symbol('a', complex=True),
|
||
}[self]
|
||
|
||
@staticmethod
|
||
def to_str(value: typ.Self) -> type:
|
||
return {
|
||
MathType.Integer: 'ℤ',
|
||
MathType.Rational: 'ℚ',
|
||
MathType.Real: 'ℝ',
|
||
MathType.Complex: 'ℂ',
|
||
}[value]
|
||
|
||
@property
|
||
def label_pretty(self) -> str:
|
||
return MathType.to_str(self)
|
||
|
||
@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,
|
||
)
|
||
|
||
|
||
####################
|
||
# - Size: 1D
|
||
####################
|
||
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 has_shape(shape: tuple[int, ...] | None):
|
||
return shape in [None, (2,), (3,), (4,), (2, 1), (3, 1), (4, 1)]
|
||
|
||
def supports_shape(self, shape: tuple[int, ...] | None):
|
||
NS = NumberSize1D
|
||
match self:
|
||
case NS.Scalar:
|
||
return shape is None
|
||
case NS.Vec2:
|
||
return shape in ((2,), (2, 1))
|
||
case NS.Vec3:
|
||
return shape in ((3,), (3, 1))
|
||
case NS.Vec4:
|
||
return shape in ((4,), (4, 1))
|
||
|
||
@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.Vec4,
|
||
(2, 1): NS.Vec2,
|
||
(3, 1): NS.Vec3,
|
||
(4, 1): NS.Vec4,
|
||
}[shape]
|
||
|
||
@property
|
||
def rows(self):
|
||
NS = NumberSize1D
|
||
return {
|
||
NS.Scalar: 1,
|
||
NS.Vec2: 2,
|
||
NS.Vec3: 3,
|
||
NS.Vec4: 4,
|
||
}[self]
|
||
|
||
@property
|
||
def cols(self):
|
||
return 1
|
||
|
||
@property
|
||
def shape(self):
|
||
NS = NumberSize1D
|
||
return {
|
||
NS.Scalar: None,
|
||
NS.Vec2: (2,),
|
||
NS.Vec3: (3,),
|
||
NS.Vec4: (4,),
|
||
}[self]
|
||
|
||
|
||
def symbol_range(sym: sp.Symbol) -> str:
|
||
return f'{sym.name} ∈ ' + (
|
||
'ℂ'
|
||
if sym.is_complex
|
||
else ('ℝ' if sym.is_real else ('ℤ' if sym.is_integer else '?'))
|
||
)
|
||
|
||
|
||
####################
|
||
# - Symbol Sizes
|
||
####################
|
||
class SimpleSize2D(enum.StrEnum):
|
||
"""Simple subset of sizes for rank-2 tensors."""
|
||
|
||
Scalar = enum.auto()
|
||
|
||
# Vectors
|
||
Vec2 = enum.auto() ## 2x1
|
||
Vec3 = enum.auto() ## 3x1
|
||
Vec4 = enum.auto() ## 4x1
|
||
|
||
# Covectors
|
||
CoVec2 = enum.auto() ## 1x2
|
||
CoVec3 = enum.auto() ## 1x3
|
||
CoVec4 = enum.auto() ## 1x4
|
||
|
||
# Square Matrices
|
||
Mat22 = enum.auto() ## 2x2
|
||
Mat33 = enum.auto() ## 3x3
|
||
Mat44 = enum.auto() ## 4x4
|
||
|
||
|
||
####################
|
||
# - 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)}
|
||
|
||
UNIT_TO_1: dict[spu.Quantity, 1] = {unit: 1 for unit in UNIT_BY_SYMBOL.values()}
|
||
|
||
|
||
####################
|
||
# - 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 <https://docs.sympy.org/latest/guides/assumptions.html#predicates>
|
||
|
||
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._
|
||
"""
|
||
## TODO: A bool flag property that does a lot of find/replace to make it super pretty
|
||
return _SYMPY_EXPR_PRINTER_STR.doprint(sp_obj)
|
||
|
||
|
||
def pretty_symbol(sym: sp.Symbol) -> str:
|
||
return f'{sym.name} ∈ ' + (
|
||
'ℤ'
|
||
if sym.is_integer
|
||
else ('ℝ' if sym.is_real else ('ℂ' if sym.is_complex 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()`.
|
||
"""
|
||
unitless_expr = spu.convert_to(sp_obj, unit) / unit if unit is not None else sp_obj
|
||
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)
|
||
|
||
|
||
@functools.cache
|
||
def unit_str_to_unit(unit_str: str) -> Unit | None:
|
||
# Edge Case: Manually Parse Degrees
|
||
## -> sp.sympify('degree') actually produces the sp.degree() function.
|
||
## -> Therefore, we must special case this particular unit.
|
||
if unit_str == 'degree':
|
||
expr = spu.degree
|
||
else:
|
||
expr = sp.sympify(unit_str).subs(UNIT_BY_SYMBOL)
|
||
|
||
if expr.has(spu.Quantity):
|
||
return expr
|
||
|
||
msg = f'No valid unit for unit string {unit_str}'
|
||
raise ValueError(msg)
|
||
|
||
|
||
####################
|
||
# - "Physical" Type
|
||
####################
|
||
def unit_dim_to_unit_dim_deps(
|
||
unit_dims: SympyType,
|
||
) -> dict[spu.dimensions.Dimension, int] | None:
|
||
dimsys_SI = spu.systems.si.dimsys_SI
|
||
|
||
# Retrieve Dimensional Dependencies
|
||
try:
|
||
return dimsys_SI.get_dimensional_dependencies(unit_dims)
|
||
|
||
# Catch TypeError
|
||
## -> Happens if `+` or `-` is in `unit`.
|
||
## -> Generally, it doesn't make sense to add/subtract differing unit dims.
|
||
## -> Thus, when trying to figure out the unit dimension, there isn't one.
|
||
except TypeError:
|
||
return None
|
||
|
||
|
||
def unit_to_unit_dim_deps(
|
||
unit: SympyType,
|
||
) -> dict[spu.dimensions.Dimension, int] | None:
|
||
# Retrieve Dimensional Dependencies
|
||
## -> NOTE: .subs() alone seems to produce sp.Symbol atoms.
|
||
## -> This is extremely problematic; `Dims` arithmetic has key properties.
|
||
## -> So we have to go all the way to the dimensional dependencies.
|
||
## -> This isn't really respecting the args, but it seems to work :)
|
||
return unit_dim_to_unit_dim_deps(
|
||
unit.subs({arg: arg.dimension for arg in unit.atoms(spu.Quantity)})
|
||
)
|
||
|
||
|
||
def compare_unit_dims(unit_dim_l: SympyType, unit_dim_r: SympyType) -> bool:
|
||
return unit_dim_to_unit_dim_deps(unit_dim_l) == unit_dim_to_unit_dim_deps(
|
||
unit_dim_r
|
||
)
|
||
|
||
|
||
def compare_unit_dim_to_unit_dim_deps(
|
||
unit_dim: SympyType, unit_dim_deps: dict[spu.dimensions.Dimension, int]
|
||
) -> bool:
|
||
return unit_dim_to_unit_dim_deps(unit_dim) == unit_dim_deps
|
||
|
||
|
||
class PhysicalType(enum.StrEnum):
|
||
"""Type identifiers for expressions with both `MathType` and a unit, aka a "physical" type."""
|
||
|
||
# Unitless
|
||
NonPhysical = enum.auto()
|
||
|
||
# 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()
|
||
Illuminance = enum.auto()
|
||
|
||
@functools.cached_property
|
||
def unit_dim(self) -> SympyType:
|
||
PT = PhysicalType
|
||
return {
|
||
PT.NonPhysical: None,
|
||
# Global
|
||
PT.Time: Dims.time,
|
||
PT.Angle: Dims.angle,
|
||
PT.SolidAngle: spu.steradian.dimension, ## 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 * spu.steradian.dimension,
|
||
PT.Illuminance: Dims.luminous_intensity / Dims.length**2,
|
||
}[self]
|
||
|
||
@staticproperty
|
||
def unit_dims() -> dict[typ.Self, SympyType]:
|
||
return {
|
||
physical_type: physical_type.unit_dim
|
||
for physical_type in list(PhysicalType)
|
||
}
|
||
|
||
@functools.cached_property
|
||
def color(self):
|
||
"""A color corresponding to the physical type.
|
||
|
||
The color selections were initially generated using AI, as this is a rote task that's better adjusted than invented.
|
||
The LLM provided the following rationale for its choices:
|
||
|
||
> Non-Physical: Grey signifies neutrality and non-physical nature.
|
||
> Global:
|
||
> Time: Blue is often associated with calmness and the passage of time.
|
||
> Angle and Solid Angle: Different shades of blue and cyan suggest angular dimensions and spatial aspects.
|
||
> Frequency and Angular Frequency: Darker shades of blue to maintain the link to time.
|
||
> Cartesian:
|
||
> Length, Area, Volume: Shades of green to represent spatial dimensions, with intensity increasing with dimension.
|
||
> Mechanical:
|
||
> Velocity and Acceleration: Red signifies motion and dynamics, with lighter reds for related quantities.
|
||
> Mass: Dark red for the fundamental property.
|
||
> Force and Pressure: Shades of red indicating intensity.
|
||
> Energy:
|
||
> Work and Power: Orange signifies energy transformation, with lighter oranges for related quantities.
|
||
> Temperature: Yellow for heat.
|
||
> Electrodynamics:
|
||
> Current and related quantities: Cyan shades indicating flow.
|
||
> Voltage, Capacitance: Greenish and blueish cyan for electrical potential.
|
||
> Impedance, Conductance, Conductivity: Purples and magentas to signify resistance and conductance.
|
||
> Magnetic properties: Magenta shades for magnetism.
|
||
> Electric Field: Light blue.
|
||
> Magnetic Field: Grey, as it can be considered neutral in terms of direction.
|
||
> Luminal:
|
||
> Luminous properties: Yellows to signify light and illumination.
|
||
>
|
||
> This color mapping helps maintain intuitive connections for users interacting with these physical types.
|
||
"""
|
||
PT = PhysicalType
|
||
return {
|
||
PT.NonPhysical: (0.75, 0.75, 0.75, 1.0), # Light Grey: Non-physical
|
||
# Global
|
||
PT.Time: (0.5, 0.5, 1.0, 1.0), # Light Blue: Time
|
||
PT.Angle: (0.5, 0.75, 1.0, 1.0), # Light Blue: Angle
|
||
PT.SolidAngle: (0.5, 0.75, 0.75, 1.0), # Light Cyan: Solid Angle
|
||
PT.Freq: (0.5, 0.5, 0.9, 1.0), # Light Blue: Frequency
|
||
PT.AngFreq: (0.5, 0.5, 0.8, 1.0), # Light Blue: Angular Frequency
|
||
# Cartesian
|
||
PT.Length: (0.5, 1.0, 0.5, 1.0), # Light Green: Length
|
||
PT.Area: (0.6, 1.0, 0.6, 1.0), # Light Green: Area
|
||
PT.Volume: (0.7, 1.0, 0.7, 1.0), # Light Green: Volume
|
||
# Mechanical
|
||
PT.Vel: (1.0, 0.5, 0.5, 1.0), # Light Red: Velocity
|
||
PT.Accel: (1.0, 0.6, 0.6, 1.0), # Light Red: Acceleration
|
||
PT.Mass: (0.75, 0.5, 0.5, 1.0), # Light Red: Mass
|
||
PT.Force: (0.9, 0.5, 0.5, 1.0), # Light Red: Force
|
||
PT.Pressure: (1.0, 0.7, 0.7, 1.0), # Light Red: Pressure
|
||
# Energy
|
||
PT.Work: (1.0, 0.75, 0.5, 1.0), # Light Orange: Work
|
||
PT.Power: (1.0, 0.85, 0.5, 1.0), # Light Orange: Power
|
||
PT.PowerFlux: (1.0, 0.8, 0.6, 1.0), # Light Orange: Power Flux
|
||
PT.Temp: (1.0, 1.0, 0.5, 1.0), # Light Yellow: Temperature
|
||
# Electrodynamics
|
||
PT.Current: (0.5, 1.0, 1.0, 1.0), # Light Cyan: Current
|
||
PT.CurrentDensity: (0.5, 0.9, 0.9, 1.0), # Light Cyan: Current Density
|
||
PT.Charge: (0.5, 0.85, 0.85, 1.0), # Light Cyan: Charge
|
||
PT.Voltage: (0.5, 1.0, 0.75, 1.0), # Light Greenish Cyan: Voltage
|
||
PT.Capacitance: (0.5, 0.75, 1.0, 1.0), # Light Blueish Cyan: Capacitance
|
||
PT.Impedance: (0.6, 0.5, 0.75, 1.0), # Light Purple: Impedance
|
||
PT.Conductance: (0.7, 0.5, 0.8, 1.0), # Light Purple: Conductance
|
||
PT.Conductivity: (0.8, 0.5, 0.9, 1.0), # Light Purple: Conductivity
|
||
PT.MFlux: (0.75, 0.5, 0.75, 1.0), # Light Magenta: Magnetic Flux
|
||
PT.MFluxDensity: (
|
||
0.85,
|
||
0.5,
|
||
0.85,
|
||
1.0,
|
||
), # Light Magenta: Magnetic Flux Density
|
||
PT.Inductance: (0.8, 0.5, 0.8, 1.0), # Light Magenta: Inductance
|
||
PT.EField: (0.75, 0.75, 1.0, 1.0), # Light Blue: Electric Field
|
||
PT.HField: (0.75, 0.75, 0.75, 1.0), # Light Grey: Magnetic Field
|
||
# Luminal
|
||
PT.LumIntensity: (1.0, 0.95, 0.5, 1.0), # Light Yellow: Luminous Intensity
|
||
PT.LumFlux: (1.0, 0.95, 0.6, 1.0), # Light Yellow: Luminous Flux
|
||
PT.Illuminance: (1.0, 1.0, 0.75, 1.0), # Pale Yellow: Illuminance
|
||
}[self]
|
||
|
||
@functools.cached_property
|
||
def default_unit(self) -> list[Unit]:
|
||
PT = PhysicalType
|
||
return {
|
||
PT.NonPhysical: None,
|
||
# 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,
|
||
}[self]
|
||
|
||
@functools.cached_property
|
||
def valid_units(self) -> list[Unit]:
|
||
"""Retrieve an ordered (by subjective usefulness) list of units for this physical type.
|
||
|
||
Notes:
|
||
The order in which valid units are declared is the exact same order that UI dropdowns display them.
|
||
|
||
**Altering the order of units breaks backwards compatibility**.
|
||
"""
|
||
PT = PhysicalType
|
||
return {
|
||
PT.NonPhysical: [None],
|
||
# Global
|
||
PT.Time: [
|
||
spu.picosecond,
|
||
femtosecond,
|
||
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 := [
|
||
terahertz,
|
||
spu.hertz,
|
||
kilohertz,
|
||
megahertz,
|
||
gigahertz,
|
||
petahertz,
|
||
exahertz,
|
||
]
|
||
),
|
||
PT.AngFreq: [spu.radian * _unit for _unit in _valid_freqs],
|
||
# Cartesian
|
||
PT.Length: (
|
||
_valid_lens := [
|
||
spu.micrometer,
|
||
spu.nanometer,
|
||
spu.picometer,
|
||
spu.angstrom,
|
||
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.kilogram,
|
||
spu.electron_rest_mass,
|
||
spu.dalton,
|
||
spu.microgram,
|
||
spu.milligram,
|
||
spu.gram,
|
||
spu.metric_ton,
|
||
],
|
||
PT.Force: [
|
||
micronewton,
|
||
nanonewton,
|
||
millinewton,
|
||
spu.newton,
|
||
spu.kg * spu.meter / spu.second**2,
|
||
],
|
||
PT.Pressure: [
|
||
spu.bar,
|
||
millibar,
|
||
spu.pascal,
|
||
hectopascal,
|
||
spu.atmosphere,
|
||
spu.psi,
|
||
spu.mmHg,
|
||
spu.torr,
|
||
],
|
||
# Energy
|
||
PT.Work: [
|
||
spu.joule,
|
||
spu.electronvolt,
|
||
],
|
||
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,
|
||
],
|
||
}[self]
|
||
|
||
@staticmethod
|
||
def from_unit(unit: Unit | None, optional: bool = False) -> typ.Self | None:
|
||
"""Attempt to determine a matching `PhysicalType` from a unit.
|
||
|
||
NOTE: It is not guaranteed that `unit` is within `valid_units`, only that it can be converted to any unit in `valid_units`.
|
||
|
||
Returns:
|
||
The matched `PhysicalType`.
|
||
|
||
If none could be matched, then either return `None` (if `optional` is set) or error.
|
||
|
||
Raises:
|
||
ValueError: If no `PhysicalType` could be matched, and `optional` is `False`.
|
||
"""
|
||
if unit is None:
|
||
return ct.PhysicalType.NonPhysical
|
||
|
||
unit_dim_deps = unit_to_unit_dim_deps(unit)
|
||
if unit_dim_deps is not None:
|
||
for physical_type, candidate_unit_dim in PhysicalType.unit_dims.items():
|
||
if compare_unit_dim_to_unit_dim_deps(candidate_unit_dim, unit_dim_deps):
|
||
return physical_type
|
||
|
||
if optional:
|
||
return None
|
||
msg = f'Could not determine PhysicalType for {unit}'
|
||
raise ValueError(msg)
|
||
|
||
@staticmethod
|
||
def from_unit_dim(
|
||
unit_dim: SympyType | None, optional: bool = False
|
||
) -> typ.Self | None:
|
||
"""Attempts to match an arbitrary unit dimension expression to a corresponding `PhysicalType`.
|
||
|
||
For comparing arbitrary unit dimensions (via expressions of `spu.dimensions.Dimension`), it is critical that equivalent dimensions are also compared as equal (ex. `mass*length/time^2 == force`).
|
||
To do so, we employ the `SI` unit conventions, for extracting the fundamental dimensional dependencies of unit dimension expressions.
|
||
|
||
Returns:
|
||
The matched `PhysicalType`.
|
||
|
||
If none could be matched, then either return `None` (if `optional` is set) or error.
|
||
|
||
Raises:
|
||
ValueError: If no `PhysicalType` could be matched, and `optional` is `False`.
|
||
"""
|
||
for physical_type, candidate_unit_dim in PhysicalType.unit_dims.items():
|
||
if compare_unit_dims(unit_dim, candidate_unit_dim):
|
||
return physical_type
|
||
|
||
if optional:
|
||
return None
|
||
msg = f'Could not determine PhysicalType for {unit_dim}'
|
||
raise ValueError(msg)
|
||
|
||
@functools.cached_property
|
||
def valid_shapes(self) -> list[typ.Literal[(3,), (2,)] | None]:
|
||
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,)],
|
||
}
|
||
|
||
return overrides.get(self, [None])
|
||
|
||
@functools.cached_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: <https://iopscience.iop.org/article/10.1088/1361-6455/aac787>
|
||
- **Conductance**: The imaginary part represents the extinction, in the Drude-model sense.
|
||
|
||
"""
|
||
MT = MathType
|
||
PT = PhysicalType
|
||
overrides = {
|
||
PT.NonPhysical: list(MT), ## Support All
|
||
# 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
|
||
}
|
||
|
||
return overrides.get(self, [MT.Real])
|
||
|
||
@staticmethod
|
||
def to_name(value: typ.Self) -> str:
|
||
if value is PhysicalType.NonPhysical:
|
||
return 'Unitless'
|
||
return PhysicalType(value).name
|
||
|
||
@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 = {
|
||
_PT.NonPhysical: None,
|
||
# 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.Voltage: spu.volt,
|
||
_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,
|
||
}
|
||
|
||
|
||
####################
|
||
# - 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 strip_unit_system(
|
||
sp_obj: SympyExpr, unit_system: UnitSystem | None = None
|
||
) -> SympyExpr:
|
||
"""Strip units occurring in the given unit system from the expression.
|
||
|
||
Unit stripping is a "dumb" operation: "Substitute any `sympy` object in `unit_system.values()` with `1`".
|
||
Obviously, the semantic correctness of this operation depends entirely on _the units adding no semantic meaning to the expression_.
|
||
|
||
Notes:
|
||
You should probably use `scale_to_unit_system()` or `convert_to_unit_system()`.
|
||
"""
|
||
if unit_system is None:
|
||
return sp_obj.subs(UNIT_TO_1)
|
||
return sp_obj.subs({unit: 1 for unit in unit_system.values() if unit is not None})
|
||
|
||
|
||
def convert_to_unit_system(
|
||
sp_obj: SympyExpr, unit_system: UnitSystem | None
|
||
) -> SympyExpr:
|
||
"""Convert an expression to the units of a given unit system."""
|
||
if unit_system is None:
|
||
return sp_obj
|
||
|
||
return spu.convert_to(
|
||
sp_obj,
|
||
{unit_system[PhysicalType.from_unit(unit)] for unit in get_units(sp_obj)},
|
||
)
|
||
|
||
|
||
def scale_to_unit_system(
|
||
sp_obj: SympyExpr, unit_system: UnitSystem | None, 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.
|
||
|
||
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(
|
||
strip_unit_system(convert_to_unit_system(sp_obj, unit_system), unit_system),
|
||
use_jax_array=use_jax_array,
|
||
)
|