Source code for Karana.Math.Kquantities

# Copyright (c) 2024-2026 Karana Dynamics Pty Ltd. All rights reserved.
#
# NOTICE TO USER:
#
# This source code and/or documentation (the "Licensed Materials") is
# the confidential and proprietary information of Karana Dynamics Inc.
# Use of these Licensed Materials is governed by the terms and conditions
# of a separate software license agreement between Karana Dynamics and the
# Licensee ("License Agreement"). Unless expressly permitted under that
# agreement, any reproduction, modification, distribution, or disclosure
# of the Licensed Materials, in whole or in part, to any third party
# without the prior written consent of Karana Dynamics is strictly prohibited.
#
# THE LICENSED MATERIALS ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND.
# KARANA DYNAMICS DISCLAIMS ALL WARRANTIES, EXPRESS OR IMPLIED, INCLUDING
# BUT NOT LIMITED TO WARRANTIES OF MERCHANTABILITY, NON-INFRINGEMENT, AND
# FITNESS FOR A PARTICULAR PURPOSE.
#
# IN NO EVENT SHALL KARANA DYNAMICS BE LIABLE FOR ANY DAMAGES WHATSOEVER,
# INCLUDING BUT NOT LIMITED TO LOSS OF PROFITS, DATA, OR USE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGES, WHETHER IN CONTRACT, TORT,
# OR OTHERWISE ARISING OUT OF OR IN CONNECTION WITH THE LICENSED MATERIALS.
#
# U.S. Government End Users: The Licensed Materials are a "commercial item"
# as defined at 48 C.F.R. 2.101, and are provided to the U.S. Government
# only as a commercial end item under the terms of this license.
#
# Any use of the Licensed Materials in individual or commercial software must
# include, in the user documentation and internal source code comments,
# this Notice, Disclaimer, and U.S. Government Use Provision.

"""Simulation quantities and unit conversion tools.

This module contains common quantities used throughout simulations. It also contains
helper functions to convert a given quantity to the simulation's units. In addition,
it contains functions to define and set a units system. The default is SI.
"""

import numpy as np
from numpy.typing import NDArray
from typing import overload
from pint import UnitRegistry, Quantity
from pint.util import UnitsContainer
from pint.facets.plain.objects import PlainUnit

# The UnitRegistry used for quantities
_ur = UnitRegistry()

# Flag to track whether ureg was ever imported
_ur_retrieved = False




def __getattr__(name):
    global _ur_retrieved
    if name == "ureg":
        # Set the flag to True. This allows us to track whether ureg was ever imported
        _ur_retrieved = True
        return _ur

    # Raise an AttributeError for other non-existent attributes
    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")





def setUnitRegistry(ur: UnitRegistry):
    """Set the UnitRegistry.

    This sets the units system used for simulation. If this is never called, the default will
    be SI units. If you want something other than SI for your simulation, you should call
    this function before ureg is ever imported from Kquantities.

    Parameters
    ----------
    ur : UnitRegistry
        The new UnitRegistry to use.
    """
    global _ur, _ur_retrieved
    _ur = ur

    if _ur_retrieved:
        from Karana.Core import warn

        warn(
            "The unit registry was set after it was already retrieved elsewhere. If units were created elsewhere, they will be associated with a different unit registry, which will lead to mismatched unit issues. Please set the unit registry first before doing anything else, even importing the unit registry from this module."
        )

    # Reset this value to false
    _ur_retrieved = False



# This variables define the known quantities.
length = _ur.m.dimensionality
angle = _ur.rad.dimensionality
mass = _ur.kg.dimensionality
time = _ur.s.dimensionality
velocity = (_ur.m / _ur.s).dimensionality
angular_velocity = (_ur.rad / _ur.s).dimensionality
acceleration = (_ur.m / _ur.s**2).dimensionality
force = _ur.N.dimensionality
torque = (_ur.N * _ur.m).dimensionality
inertia = (_ur.kg * _ur.m**2).dimensionality
gravitational_parameter = (_ur.m**3 / _ur.s**2).dimensionality
linear_spring_constant = (_ur.N / _ur.m).dimensionality
linear_damping_constant = (_ur.N * _ur.s / _ur.m).dimensionality
rotational_spring_constant = (_ur.N * _ur.m / _ur.rad).dimensionality
rotational_damping_constant = (_ur.N * _ur.m * _ur.s / _ur.rad).dimensionality




def getDefaultUnits(quantity: UnitsContainer) -> PlainUnit:
    """Get the default units for a quantity.

    Parameters
    ----------
    quantity : UnitsContainer
        The quantity to get default units for.

    Returns
    -------
    PlainUnit
        The default units for the provided quantity.
    """
    # Set ur_retrieved to True, as this uses the UnitRegistry
    global _ur_retrieved
    _ur_retrieved = True

    if quantity == angle:
        return _ur.get_base_units(_ur.rad)[1]

    dark = 1.0
    for k, v in quantity.items():
        dark *= _ur.get_base_units(next(iter(_ur.get_compatible_units(k))))[1] ** v

    return dark.u



@overload
def convert(v: float) -> float: ...
@overload
def convert(v: NDArray[np.float64]) -> NDArray[np.float64]: ...
@overload
def convert(v: Quantity) -> Quantity: ...


def convert(v: object) -> object:
    """Convert this quantity to the current units system."""
    # If v is not a quantity, i.e., has no units, then there
    # is no need to convert it.
    if not isinstance(v, Quantity):
        return v
    return v.to_base_units()