cuequivariance

cuEquivariance コア

cuequivariance（import cuequivariance as cue）を使用して、群表現とテンソル縮約を記述します。要求された演算を表現できる場合は、組み込みの記述子を優先してください。カスタム縮約の場合に限り、SegmentedTensorProduct を構築します。コア多項式は NumPy を使って CPU で評価できます。学習可能なフレームワークパラメーターを作成したり、GPU カーネルを選択したりすることはありません。

ワークフローを選択する

• 組み込みの群と密な記述子には、以下の例を使用します。
• 新しい群については、カスタム既約表現を参照してください。
• カスタム縮約、CG パス、多項式の確認については、セグメント化多項式を参照してください。
• 既約表現をキーとする辞書については、ir_dict のワークフローを参照してください。CPU の例と、JAX 実行との境界について説明しています。

説明や実行可能なコードの提示を求められた場合は、最も近い例を応用してください。回答するためにパッケージをインストールする必要はありません。選択したワークフローのリファレンスを参照してください。ユーザーから実行を求められた場合、または統合や正確性に関する質問で実行時の確認が必要な場合は、コードを実行してください。インストール済みの API が異なる場合や数値チェックに失敗した場合は、上流のソースを確認してください。

実行環境を確認する

これらの例は cuEquivariance 0.12.0 で確認済みです。実行する際は、選択したインタープリターと、そのワークフローに必要な API を合わせて確認してください。

from importlib.metadata import PackageNotFoundError, version
import sys

print("Python:", sys.executable)
try:
    installed = version("cuequivariance")
except PackageNotFoundError:
    print("cuequivariance is not installed in this interpreter")
else:
    import cuequivariance as cue

    print("cuEquivariance:", installed)
    print("IrDictPolynomial available:", hasattr(cue, "IrDictPolynomial"))

互換性のあるプロジェクト環境と、その依存関係の制約を再利用してください。実行に不足しているパッケージが必要な場合は、分離された環境での例の実行を使用してください。既存の仮想環境とシステム／ユーザーの Python インストールは保持し、環境を削除したり、セットアップを修復するためにパッケージマネージャーの保護機能を回避したりしないでください。インストールできない、または禁止されている場合は、実行可能なコードを提示し、実行していないチェックを明記してください。コアの例に必要なのは NumPy だけで、CUDA も JAX／PyTorch も必要ありません。API が利用できない場合は、タスクを満たすならサポートされている密な記述子を使用してください。メソッドをでっち上げたり、環境のアップグレードを強制したりしないでください。

フレームワークに関するタスクでは、インストール済みの cuequivariance_jax または cuequivariance_torch のシグネチャと、そのバージョンのバックエンド要件を確認してください。 公式ドキュメントと 0.12.0 のソースでは、これらの別個のパッケージについて説明しています。

群、重複度、レイアウト

群	構築方法	次元と意味
SO(3)	cue.SO3(l)	2*l + 1；回転；非負の整数 l
O(3)	cue.O3(l, p)	2*l + 1；パリティ p=+1 または -1 も指定します。例：0e、1o
SU(2)	cue.SU2(j)	2*j + 1；非負の整数または半整数のスピン

cue.Irreps("SO3", "16x0 + 4x1 + 2x2") の次元は 38 です。これを反復すると、(multiplicity, irrep) ブロックが得られます。Irreps に整数を掛けると重複度がスケーリングされます。個々の既約表現に対する ir1 * ir2 は選択則を表します。SO(3) では、許される出力の角運動量は abs(l1-l2) から l1+l2 までです。O(3) では、出力のパリティも p1*p2 と一致する必要があります。

レイアウトはベクトル全体ではなく、各既約表現ブロック内で定義されます。

レイアウト	平坦化前のブロック形状
cue.ir_mul	(..., ir.dim, mul)
cue.mul_ir	(..., mul, ir.dim)

記述子から e.inputs[i].layout と e.outputs[i].layout を読み取ってください。ここでのテンソル積の例では ir_mul を使用しています。多くの PyTorch レイヤーは既定で mul_ir を使用し、明示的なレイアウトオプションを受け付けます。JAX の辞書機能では mul_ir がよく使われますが、多項式の実行では ir_mul が必要になる場合があります。平坦化の前に、実際のデータブロックをブロックごとに転置して変換してください。メタデータの変更や reshape の呼び出しだけでは、要素の順序は変わりません。同様に、既約表現のメタデータをソートまたは簡略化しても、既存のテンソルの順序は変わりません。評価時は、記述子のブロック順序を維持してください。

実行可能な密テンソル積

ユーザーからすべての出力既約表現を求められたら、full_tensor_product から始めます。 入力は2つで、重みはありません。この例には、すべての入力定義が含まれています。

import numpy as np
import cuequivariance as cue

irreps1 = cue.Irreps("SO3", "2x0 + 1x1 + 1x2")
irreps2 = cue.Irreps("SO3", "1x0 + 1x1")
e = cue.descriptors.full_tensor_product(irreps1, irreps2)
rng = np.random.default_rng(0)
x1 = rng.normal(size=(4, e.inputs[0].dim))
x2 = rng.normal(size=(4, e.inputs[1].dim))
[out] = e(x1, x2)
assert out.shape == (4, 40)
assert np.isfinite(out).all()

# A multiplicity summary for display; keep e.outputs[0] for out's actual layout.
summary = e.outputs[0].irreps.sort().irreps.simplify()
assert summary == cue.Irreps("SO3", "3x0 + 5x1 + 3x2 + 1x3")

記述子を選択する

以下の名前はすべて cue.descriptors にあります。重みを含め、配列を割り当てるには、返された入力／出力のメタデータを調べてください。重みの数を推測しないでください。

記述子	用途とオペランドの順序
linear(irreps_in, irreps_out)	一致する既約表現の重複度を混合します。重み、入力 → 出力
fully_connected_tensor_product(a, b, out)	出力の重複度を正確に指定します。重み、入力1、入力2 → 出力
full_tensor_product(a, b)	許容されるすべての結合。重みなし。入力1、入力2 → 出力
channelwise_tensor_product(a, b, irreps3_filter=...)	チャネルの組に重みを付けます。各経路の出力重複度は mul1*mul2 です
elementwise_tensor_product(a, b)	重みなしでチャネルをペアにします。入力間で既約表現のコピー数が等しく、ブロック境界に互換性がある必要があります
spherical_harmonics(cue.O3(1, -1), [0, 1, 2, 3])	ベクトル → パリティを含む調和関数。単位球面上の値が必要な場合は、入力ベクトルを別途正規化してください
symmetric_contraction(irreps_in, irreps_out, (1, 2, 3))	MACE 形式の多項式相関。記述子の形状と重複度の制約を確認してください

チャネル単位積の3番目の引数は、既約表現の型をフィルタリングします。出力の重複度を設定するものではありません。チャネル単位積ではチャネルの組（u, v）が保持されますが、要素単位積では対応するチャネル（u）が保持されます。密なチャネル単位記述子で重複する出力既約表現を統合するには、simplify_irreps3=True を使用します。

EquivariantPolynomial は、inputs と outputs に密な表現メタデータを格納し、polynomial に縮約を格納します。呼び出すと、出力が1つの場合でも NumPy 出力のリストが返されます。対応する _ir_dict 版は、既約表現ごとに分割済みの多項式とメタデータを返します。リンク先のワークフローを参照してください。

独自の実装では、出力次元と有限値を確認してから、非自明な群変換に対する共変性を検証してください。入力と期待される出力の両方で、同じ群、基底、ブロック順序、レイアウトを使用してください。