Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Backport PR #350 on branch versions/v1.0.x (🎨 style(quantity): better layout for constructors) #351

Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
228 changes: 115 additions & 113 deletions src/unxt/_src/quantity/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,123 +113,15 @@ class AbstractQuantity(
# Constructors

@classmethod
@dispatch
@dispatch.abstract # type: ignore[misc]
def from_(
cls: "type[AbstractQuantity]",
value: ArrayLike | ArrayLikeSequence,
unit: Any,
/,
*,
dtype: Any = None,
) -> "AbstractQuantity":
"""Construct a `unxt.Quantity` from an array-like value and a unit.

:param value: The array-like value.
:param unit: The unit of the value.
:param dtype: The data type of the array (keyword-only).

Examples
--------
For this example we'll use the `Quantity` class. The same applies to
any subclass of `AbstractQuantity`.

>>> import jax.numpy as jnp
>>> import unxt as u

>>> x = jnp.array([1.0, 2, 3])
>>> u.Quantity.from_(x, "m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

>>> u.Quantity.from_([1.0, 2, 3], "m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

>>> u.Quantity.from_((1.0, 2, 3), "m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

"""
# Dispatch on both arguments.
# Construct using the standard `__init__` method.
return cls(jnp.asarray(value, dtype=dtype), unit)

@classmethod # type: ignore[no-redef]
@dispatch
def from_(
cls: "type[AbstractQuantity]",
value: ArrayLike | ArrayLikeSequence,
/,
*,
unit: Any,
dtype: Any = None,
*args: Any,
**kwargs: Any,
) -> "AbstractQuantity":
"""Make a `unxt.AbstractQuantity` from an array-like value and a unit kwarg.

Examples
--------
For this example we'll use the `unxt.Quantity` class. The same applies
to any subclass of `unxt.AbstractQuantity`.
raise NotImplementedError # pragma: no cover

>>> import unxt as u
>>> u.Quantity.from_([1.0, 2, 3], unit="m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

"""
# Dispatch on the `value` only. Dispatch to the full constructor.
return cls.from_(value, unit, dtype=dtype)

@classmethod # type: ignore[no-redef]
@dispatch
def from_(
cls: "type[AbstractQuantity]", *, value: Any, unit: Any, dtype: Any = None
) -> "AbstractQuantity":
"""Construct a `AbstractQuantity` from value and unit kwargs.

Examples
--------
For this example we'll use the `Quantity` class. The same applies to
any subclass of `AbstractQuantity`.

>>> import unxt as u
>>> u.Quantity.from_(value=[1.0, 2, 3], unit="m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

"""
# Dispatched on no argument. Dispatch to the full constructor.
return cls.from_(value, unit, dtype=dtype)

@classmethod # type: ignore[no-redef]
@dispatch
def from_(
cls: "type[AbstractQuantity]", mapping: Mapping[str, Any]
) -> "AbstractQuantity":
"""Construct a `Quantity` from a Mapping.

Parameters
----------
mapping : Mapping[str, Any]
Mapping of the fields of the `Quantity`, e.g. 'value' and 'unit'.

Examples
--------
For this example we'll use the `Quantity` class. The same applies to
any subclass of `AbstractQuantity`.

>>> import jax.numpy as jnp
>>> import unxt as u

>>> x = jnp.array([1.0, 2, 3])
>>> q = u.Quantity.from_({"value": x, "unit": "m"})
>>> q
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

>>> u.Quantity.from_({"value": q, "unit": "km"})
Quantity['length'](Array([0.001, 0.002, 0.003], dtype=float32), unit='km')

"""
# Dispatch on both arguments.
# Construct using the standard `__init__` method.
return cls.from_(**mapping)

# See below for additional constructors.
# See below for additional constructors

# ===============================================================
# Quantity API
Expand Down Expand Up @@ -892,6 +784,116 @@ def __repr__(self) -> str:
# Register additional constructors


@AbstractQuantity.from_.dispatch
def from_(
cls: type[AbstractQuantity],
value: ArrayLike | ArrayLikeSequence,
unit: Any,
/,
*,
dtype: Any = None,
) -> AbstractQuantity:
"""Construct a `unxt.Quantity` from an array-like value and a unit.

:param value: The array-like value.
:param unit: The unit of the value.
:param dtype: The data type of the array (keyword-only).

Examples
--------
For this example we'll use the `Quantity` class. The same applies to
any subclass of `AbstractQuantity`.

>>> import jax.numpy as jnp
>>> import unxt as u

>>> x = jnp.array([1.0, 2, 3])
>>> u.Quantity.from_(x, "m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

>>> u.Quantity.from_([1.0, 2, 3], "m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

>>> u.Quantity.from_((1.0, 2, 3), "m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

"""
# Dispatch on both arguments.
# Construct using the standard `__init__` method.
return cls(jnp.asarray(value, dtype=dtype), unit)


@AbstractQuantity.from_.dispatch
def from_(
cls: type[AbstractQuantity],
value: ArrayLike | ArrayLikeSequence,
/,
*,
unit: Any,
dtype: Any = None,
) -> AbstractQuantity:
"""Make a `unxt.AbstractQuantity` from an array-like value and a unit kwarg.

Examples
--------
For this example we'll use the `unxt.Quantity` class. The same applies
to any subclass of `unxt.AbstractQuantity`.

>>> import unxt as u
>>> u.Quantity.from_([1.0, 2, 3], unit="m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

"""
# Dispatch on the `value` only. Dispatch to the full constructor.
return cls.from_(value, unit, dtype=dtype)


@AbstractQuantity.from_.dispatch
def from_(
cls: type[AbstractQuantity], *, value: Any, unit: Any, dtype: Any = None
) -> AbstractQuantity:
"""Construct a `AbstractQuantity` from value and unit kwargs.

Examples
--------
For this example we'll use the `Quantity` class. The same applies to
any subclass of `AbstractQuantity`.

>>> import unxt as u
>>> u.Quantity.from_(value=[1.0, 2, 3], unit="m")
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

"""
# Dispatched on no argument. Dispatch to the full constructor.
return cls.from_(value, unit, dtype=dtype)


@AbstractQuantity.from_.dispatch
def from_(cls: type[AbstractQuantity], mapping: Mapping[str, Any]) -> AbstractQuantity:
"""Construct a `Quantity` from a Mapping.

Examples
--------
For this example we'll use the `Quantity` class. The same applies to
any subclass of `AbstractQuantity`.

>>> import jax.numpy as jnp
>>> import unxt as u

>>> x = jnp.array([1.0, 2, 3])
>>> q = u.Quantity.from_({"value": x, "unit": "m"})
>>> q
Quantity['length'](Array([1., 2., 3.], dtype=float32), unit='m')

>>> u.Quantity.from_({"value": q, "unit": "km"})
Quantity['length'](Array([0.001, 0.002, 0.003], dtype=float32), unit='km')

"""
# Dispatch on both arguments.
# Construct using the standard `__init__` method.
return cls.from_(**mapping)


@AbstractQuantity.from_.dispatch
def from_(
cls: type[AbstractQuantity],
Expand Down
Loading