Jet

Convenience class to work with jets (evaluation point & derivative table)

class ndmap.jet.Jet(dimension: tuple[int, ...], order: tuple[int, ...], *, initialize: bool = True, point: list[torch.Tensor] | None = None, jacobian: Callable | None = None, dtype: torch.dtype = torch.float64, device: torch.device = torch.device)[source]

Convenience class to work with jets (evaluation point & derivative table)

Return type:

Jet class instance

Parameters:

dimension (tuple[int, ...], positive) – dimensions
order (tuple[int, ...], non-negative) – maximum derivative orders
initialize (bool) – flag to initialize identity derivative table, optional, default=True
point (Optional[Point]) – evaluation point, default=None
jacobian (Optional[Callable]) – torch.func.jacfwd (default) or torch.func.jacrev
dtype (torch.dtype, default=torch.float64) – data type
device (torch.device, default=torch.device('cpu')) – data device

dimension

dimensions

Type:: tuple[int, …], positive

order

maximum derivative orders

Type:: tuple[int, …], non-negative

initialize

flag to initialize identity derivative table, optional, default=True

Type:: bool

point

evaluation point, default=None

Type:: Point

jacobian

torch.func.jacfwd (default) or torch.func.jacrev

Type:: Optional[Callable]

dtype

data type

Type:: torch.dtype, default=torch.float64

device

data device

Type:: torch.device, default=torch.device(‘cpu’)

state

state

Type:: State

knobs

knobs

Type:: Knobs

table

table representation

Type:: Table

series

series representation

Type:: Series

signature

derivative table elements bottom elements signatures

Type:: list[tuple[int, …]]

parametetric

parametric table

Type:: Table

compliant(other: Jet) → bool[source]

Check jets are compliant (can be composed)

Parameters:: other (Jet) – other jet
Return type:: bool

Examples

>>> j1 = Jet((2, 2), (1, 1))
>>> j2 = Jet((2, 2), (2, 1))
>>> j3 = Jet((2, 2), (1, 1))
>>> j1.compliant(j2)
False
>>> j1.compliant(j3)
True

compose(other: Jet) → Jet[source]

Compose jets (evaluate other jet at self jet)

Parameters:: other (Jet) – other jet
Return type:: Jet

Examples

>>> j1 = Jet((2, 2), (1, 1), dtype=torch.float32)
>>> j2 = Jet((2, 2), (1, 1), dtype=torch.float32)
>>> j1.compose(j2).table
[[tensor([0., 0.]),
tensor([[0., 0.],
        [0., 0.]])],
[tensor([[1., 0.],
        [0., 1.]]),
tensor([[[0., 0.],
        [0., 0.]],

[[0., 0.], [0., 0.]]])]]

evaluate(delta: list[torch.Tensor]) → torch.Tensor[source]

Evaluate jet derivative table at a given delta deviation

Parameters:: delta (Delta) – delta deviation
Return type:: Tensor

Examples

>>> import torch
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> j = Jet((2, 2), (2, 1), initialize=True, point=[x, k], dtype=torch.float32)
>>> dx = torch.tensor([1.0, 1.0])
>>> dk = torch.tensor([1.0, 1.0])
>>> j.evaluate([dx, dk])
tensor([1., 1.])

classmethod from_mapping(dimension: tuple[int, ...], order: tuple[int, ...], point: list[torch.Tensor], function: Callable, *args: tuple, jacobian: Callable | None = None, dtype: torch.dtype = torch.float64, device: torch.device = torch.device) → Jet[source]

Jet initialization from mapping

Parameters:

dimension (tuple[int, ...], positive) – dimensions
order (tuple[int, ...], non-negative) – maximum derivative orders
point (Point) – evaluation point
function (Mapping) – input function
*args (tuple) – additional function arguments
jacobian (Optional[Callable]) – torch.func.jacfwd (default) or torch.func.jacrev
dtype (torch.dtype, default=torch.float64) – data type
device (torch.device, default=torch.device('cpu')) – data device

Return type:

Jet

Examples

>>> import torch
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> def mapping(x, k):
...     x1, x2 = x
...     k1, k2 = k
...     return torch.stack([x1*k1, x2*k2])
>>> j = Jet.from_mapping((2, 2), (1, 1), [x, k], mapping, dtype=torch.float32)
>>> j.table
[[tensor([0., 0.]),
tensor([[0., 0.],
        [0., 0.]])],
[tensor([[0., 0.],
        [0., 0.]]),
tensor([[[1., 0.],
        [0., 0.]],

[[0., 0.], [0., 1.]]])]]

classmethod from_series(dimension: tuple[int, ...], order: tuple[int, ...], point: list[torch.Tensor], series: dict[tuple[int, ...], torch.Tensor], jacobian: Callable | None = None, dtype: torch.dtype = torch.float64, device: torch.device = torch.device) → Jet[source]

Jet initialization from series

Parameters:

dimension (tuple[int, ...], positive) – dimensions
order (tuple[int, ...], non-negative) – maximum derivative orders
point (list[Tensor]) – evaluation point
series (Series) – input series
jacobian (Optional[Callable]) – torch.func.jacfwd (default) or torch.func.jacrev
dtype (torch.dtype, default=torch.float64) – data type
device (torch.device, default=torch.device('cpu')) – data device

Returns:

Jet
>>> import torch
>>> from ndmap.util import curry_apply
>>> from ndmap.series import series
>>> from ndmap.propagate import identity
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> def mapping(x, k)
… x1, x2 = x
… k1, k2 = k
… return torch.stack([x1*k1, x2*k2])
>>> s = [*identity((1, 1), [x, k], flag=True).keys()]
>>> s = series(s, curry_apply(mapping, (2, 2)), *x, *k)
>>> j = Jet.from_series((2, 2), (1, 1), [x, k], s)
>>> j.table
[[tensor([0., 0.]),
tensor([[0., 0.], – [0., 0.]])],
[tensor([[0., 0.], – [0., 0.]]),
tensor([[[1., 0.], – [0., 0.]],

[[0., 0.], [0., 1.]]])]]

classmethod from_table(dimension: tuple[int, ...], order: tuple[int, ...], point: list[torch.Tensor], table: list, jacobian: Callable | None = None, dtype: torch.dtype = torch.float64, device: torch.device = torch.device) → Jet[source]

Jet initialization from table

Parameters:

dimension (tuple[int, ...], positive) – dimensions
order (tuple[int, ...], non-negative) – maximum derivative orders
point (Point) – evaluation point
table (Table) – input (derivative) table
jacobian (Optional[Callable]) – torch.func.jacfwd (default) or torch.func.jacrev
dtype (torch.dtype, default=torch.float64) – data type
device (torch.device, default=torch.device('cpu')) – data device

Return type:

Jet

Examples

>>> import torch
>>> from ndmap.derivative import derivative
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> def mapping(x, k):
...     x1, x2 = x
...     k1, k2 = k
...     return torch.stack([x1*k1, x2*k2])
>>> t = derivative((1, 1), mapping, [x, k])
>>> j = Jet.from_table((2, 2), (1, 1), [x, k], t, dtype=torch.float32)
>>> j.table
[[tensor([0., 0.]),
tensor([[0., 0.],
        [0., 0.]])],
[tensor([[0., 0.],
        [0., 0.]]),
tensor([[[1., 0.],
        [0., 0.]],

[[0., 0.], [0., 1.]]])]]

property parametetric: list

Get parametric table (first subtable)

Parameters:: None
Return type:: Table

Examples

>>> import torch
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> j = Jet((2, 2), (2, 1), initialize=True, point=[x, k], dtype=torch.float32)
>>> j.parametetric
[tensor([0., 0.]),
tensor([[0., 0.],
        [0., 0.]])]

propagate(function: Callable, *pars: tuple) → Jet[source]

Propagate jet.

Parameters:

function (Mapping) – input function
knobs (Knobs) – input function knobs
*pars (tuple) – additional function arguments

Return type:

Jet

Examples

>>> import torch
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> def mapping(x, k):
...     x1, x2 = x
...     k1, k2 = k
...     return torch.stack([x1*k1, x2*k2])
>>> j = Jet((2, 2), (1, 1), point=[x, k], dtype=torch.float32)
>>> j.table
[[tensor([0., 0.]),
tensor([[0., 0.],
        [0., 0.]])],
[tensor([[1., 0.],
        [0., 1.]]),
tensor([[[0., 0.],
        [0., 0.]],

[[0., 0.], [0., 0.]]])]]

>>> j.propagate(mapping).table
[[tensor([0., 0.]),
tensor([[0., 0.],
        [0., 0.]])],
[tensor([[0., 0.],
        [0., 0.]]),
tensor([[[1., 0.],
        [0., 0.]],

[[0., 0.], [0., 1.]]])]]

property series: dict[tuple[int, ...], torch.Tensor]

Series representation

Parameters:: None
Return type:: Series

Examples

>>> import torch
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> j = Jet((2, 2), (2, 1), initialize=True, point=[x, k], dtype=torch.float32)
>>> j.series
{(0, 0, 0, 0): tensor([0., 0.]),
(0, 0, 1, 0): tensor([0., 0.]),
(0, 0, 0, 1): tensor([0., 0.]),
(1, 0, 0, 0): tensor([1., 0.]),
(0, 1, 0, 0): tensor([0., 1.]),
(1, 0, 1, 0): tensor([0., 0.]),
(1, 0, 0, 1): tensor([0., 0.]),
(0, 1, 1, 0): tensor([0., 0.]),
(0, 1, 0, 1): tensor([0., 0.]),
(2, 0, 0, 0): tensor([0., 0.]),
(1, 1, 0, 0): tensor([0., 0.]),
(0, 2, 0, 0): tensor([0., 0.]),
(2, 0, 1, 0): tensor([0., 0.]),
(2, 0, 0, 1): tensor([0., 0.]),
(1, 1, 1, 0): tensor([0., 0.]),
(1, 1, 0, 1): tensor([0., 0.]),
(0, 2, 1, 0): tensor([0., 0.]),
(0, 2, 0, 1): tensor([0., 0.])}

property signature: list[tuple[int, ...]] | list[tuple[tuple[int, ...], float]]

Compute derivative table elements bottom elements signatures

Parameters:: None
Returns:: bottom table elements signatures
Return type:: Signature

Examples

>>> import torch
>>> x = torch.tensor([0.0, 0.0])
>>> k = torch.tensor([0.0, 0.0])
>>> j = Jet((2, 2), (2, 1), initialize=True, point=[x, k], dtype=torch.float32)
>>> j.signature
[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]