lloca.reps.tensorreps_transform.TensorRepsTransform

class lloca.reps.tensorreps_transform.TensorRepsTransform(reps, use_naive=False)[source]

Bases: Module

Tensor representation transformation module.

Parameters:

  • reps (TensorReps) – Tensor representations to transform.

  • use_naive (bool) – Whether to use the naive transformation method. If False, uses an efficient transformation method. Default is False.

forward(tensor, frames)[source]

Apply a transformation to a tensor of a given representation.

Parameters:

  • tensor (torch.Tensor) – The tensor to transform, shape (…, self.reps.dim).

  • frames (Frames) – The local frames to apply the transformation with, shape (…, 4, 4).

Returns:

The transformed tensor, shape (…, self.reps.dim).

Return type:

torch.Tensor

transform_parity(tensor, frames)[source]

Parity transform: Multiply parity-odd states by sign(det Lambda).

Parameters:

  • tensor (torch.Tensor) – The tensor to transform, shape (N, self.reps.dim).

  • frames (Frames) – The local frames to apply the transformation with, shape (N, 4, 4).

Returns:

The transformed tensor, shape (N, self.reps.dim).

Return type:

torch.Tensor